20 #include <rmm/detail/error.hpp>
21 #include <rmm/detail/export.hpp>
22 #include <rmm/detail/runtime_async_alloc.hpp>
23 #include <rmm/detail/thrust_namespace.h>
27 #include <cuda/std/type_traits>
28 #include <cuda_runtime_api.h>
33 namespace RMM_NAMESPACE {
60 posix_file_descriptor = 0x1,
86 std::optional<std::size_t> release_threshold = {},
87 std::optional<allocation_handle_type> export_handle_type = {})
90 RMM_EXPECTS(rmm::detail::runtime_async_alloc::is_supported(),
91 "cudaMallocAsync not supported with this CUDA driver/runtime version");
94 cudaMemPoolProps pool_props{};
95 pool_props.allocType = cudaMemAllocationTypePinned;
96 pool_props.handleTypes =
static_cast<cudaMemAllocationHandleType
>(
97 export_handle_type.value_or(allocation_handle_type::none));
99 rmm::detail::runtime_async_alloc::is_export_handle_type_supported(pool_props.handleTypes),
100 "Requested IPC memory handle type not supported");
101 pool_props.location.type = cudaMemLocationTypeDevice;
103 cudaMemPool_t cuda_pool_handle{};
104 RMM_CUDA_TRY(cudaMemPoolCreate(&cuda_pool_handle, &pool_props));
105 pool_ = cuda_async_view_memory_resource{cuda_pool_handle};
110 int driver_version{};
111 RMM_CUDA_TRY(cudaDriverGetVersion(&driver_version));
112 constexpr
auto min_async_version{11050};
113 if (driver_version < min_async_version) {
116 cudaMemPoolSetAttribute(pool_handle(), cudaMemPoolReuseAllowOpportunistic, &disabled));
122 uint64_t threshold = release_threshold.value_or(total);
124 cudaMemPoolSetAttribute(pool_handle(), cudaMemPoolAttrReleaseThreshold, &threshold));
128 auto const pool_size = initial_pool_size.value_or(free / 2);
138 [[nodiscard]] cudaMemPool_t
pool_handle() const noexcept {
return pool_.pool_handle(); }
142 RMM_ASSERT_CUDA_SUCCESS(cudaMemPoolDestroy(pool_handle()));
144 cuda_async_memory_resource(cuda_async_memory_resource
const&) =
delete;
145 cuda_async_memory_resource(cuda_async_memory_resource&&) =
delete;
146 cuda_async_memory_resource& operator=(cuda_async_memory_resource
const&) =
delete;
147 cuda_async_memory_resource& operator=(cuda_async_memory_resource&&) =
delete;
150 cuda_async_view_memory_resource pool_{};
164 ptr = pool_.allocate(bytes, stream);
178 pool_.deallocate(ptr, bytes, stream);
188 [[nodiscard]]
bool do_is_equal(device_memory_resource
const& other)
const noexcept
override
190 auto const* async_mr =
dynamic_cast<cuda_async_memory_resource const*
>(&other);
191 return (async_mr !=
nullptr) && (this->pool_handle() == async_mr->pool_handle());
Strongly-typed non-owning wrapper for CUDA streams with default constructor.
Definition: cuda_stream_view.hpp:39
device_memory_resource derived class that uses cudaMallocAsync/cudaFreeAsync for allocation/deallocat...
Definition: cuda_async_memory_resource.hpp:45
allocation_handle_type
Flags for specifying memory allocation handle types.
Definition: cuda_async_memory_resource.hpp:58
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...
Definition: cuda_async_memory_resource.hpp:85
cudaMemPool_t pool_handle() const noexcept
Returns the underlying native handle to the CUDA pool.
Definition: cuda_async_memory_resource.hpp:138
Base class for all librmm device memory allocation.
Definition: device_memory_resource.hpp:93
std::pair< std::size_t, std::size_t > available_device_memory()
Returns the available and total device memory in bytes for the current device.
Definition: cuda_device.hpp:123
cuda_device_id get_current_cuda_device()
Returns a cuda_device_id for the current device.
Definition: cuda_device.hpp:99
static constexpr cuda_stream_view cuda_stream_default
Static cuda_stream_view of the default stream (stream 0), for convenience.
Definition: cuda_stream_view.hpp:125
constexpr value_type value() const noexcept
The wrapped integer value.
Definition: cuda_device.hpp:57