Skip to content

Commit

Permalink
[CUDAX] Rename memory resource and memory pool from async to device (N…
Browse files Browse the repository at this point in the history 
…VIDIA#2710)

* Rename the type

* Update tests

* Rename async memory pool

* Rename the tests

* Change name in the docs

* Generalise the memory_pool_properties name

* Fix docs

---------

Co-authored-by: Michael Schellenberger Costa <miscco@nvidia.com>
  • Loading branch information
@pciolkosz @miscco
pciolkosz and miscco authored Nov 6, 2024
1 parent 2864f2b commit 6edb860
Show file tree
Hide file tree
Showing 9 changed files with 200 additions and 200 deletions.
98 changes: 49 additions & 49 deletions ...l/__memory_resource/async_memory_pool.cuh → .../__memory_resource/device_memory_pool.cuh
View file
Original file line number Diff line number Diff line change
Expand Up @@ -8,8 +8,8 @@
//
//===----------------------------------------------------------------------===//

#ifndef _CUDAX__MEMORY_RESOURCE_CUDA_MEMORY_POOL
#define _CUDAX__MEMORY_RESOURCE_CUDA_MEMORY_POOL
#ifndef _CUDAX__MEMORY_RESOURCE_DEVICE_MEMORY_POOL
#define _CUDAX__MEMORY_RESOURCE_DEVICE_MEMORY_POOL

#include <cuda/std/detail/__config>

Expand Down Expand Up @@ -42,7 +42,7 @@
# if _CCCL_STD_VER >= 2014

//! @file
//! The \c async_memory_pool class provides a wrapper around a `cudaMempool_t`.
//! The \c device_memory_pool class provides a wrapper around a `cudaMempool_t`.
namespace cuda::experimental::mr
{

Expand Down Expand Up @@ -105,20 +105,20 @@ enum class cudaMemAllocationHandleType
cudaMemHandleTypeFabric = 0x8, ///< Allows a fabric handle to be used for exporting. (cudaMemFabricHandle_t)
};

//! @brief \c async_memory_pool_properties is a wrapper around properties passed to \c async_memory_pool to create a
//! @brief \c memory_pool_properties is a wrapper around properties passed to \c device_memory_pool to create a
//! <a href="https://docs.nvidia.com/cuda/cuda-runtime-api/group__CUDART__MEMORY__POOLS.html">cudaMemPool_t</a>.
struct async_memory_pool_properties
struct memory_pool_properties
{
size_t initial_pool_size = 0;
size_t release_threshold = 0;
cudaMemAllocationHandleType allocation_handle_type = cudaMemAllocationHandleType::cudaMemHandleTypeNone;
};

//! @brief \c async_memory_pool is an owning wrapper around a
//! @brief \c device_memory_pool is an owning wrapper around a
//! <a href="https://docs.nvidia.com/cuda/cuda-runtime-api/group__CUDART__MEMORY__POOLS.html">cudaMemPool_t</a>.
//!
//! It handles creation and destruction of the underlying pool utilizing the provided \c async_memory_pool_properties.
class async_memory_pool
//! It handles creation and destruction of the underlying pool utilizing the provided \c memory_pool_properties.
class device_memory_pool
{
private:
::cudaMemPool_t __pool_handle_ = nullptr;
Expand Down Expand Up @@ -164,10 +164,10 @@ private:
//! @throws cuda_error If the creation of the CUDA memory pool failed.
//! @returns The created CUDA memory pool.
_CCCL_NODISCARD static cudaMemPool_t
__create_cuda_mempool(const int __device_id, async_memory_pool_properties __properties) noexcept
__create_cuda_mempool(const int __device_id, memory_pool_properties __properties) noexcept
{
::cuda::experimental::mr::__device_supports_stream_ordered_allocations(__device_id);
async_memory_pool::__cuda_supports_export_handle_type(__device_id, __properties.allocation_handle_type);
device_memory_pool::__cuda_supports_export_handle_type(__device_id, __properties.allocation_handle_type);

::cudaMemPoolProps __pool_properties{};
__pool_properties.allocType = ::cudaMemAllocationTypePinned;
Expand Down Expand Up @@ -210,50 +210,50 @@ private:
void* __ptr{nullptr};
_CCCL_TRY_CUDA_API(
::cudaMallocAsync,
"async_memory_pool failed to allocate the initial pool size",
"device_memory_pool failed to allocate the initial pool size",
&__ptr,
__properties.initial_pool_size,
__temp_stream.get());

_CCCL_ASSERT_CUDA_API(
::cudaFreeAsync, "async_memory_pool failed to free the initial pool allocation", __ptr, __temp_stream.get());
::cudaFreeAsync, "device_memory_pool failed to free the initial pool allocation", __ptr, __temp_stream.get());
}
return __cuda_pool_handle;
}

struct __from_handle_t
{};

//! @brief Constructs a \c async_memory_pool from a handle taking ownership of the pool
//! @brief Constructs a \c device_memory_pool from a handle taking ownership of the pool
//! @param __handle The handle to the existing pool
explicit async_memory_pool(__from_handle_t, ::cudaMemPool_t __handle) noexcept
explicit device_memory_pool(__from_handle_t, ::cudaMemPool_t __handle) noexcept
: __pool_handle_(__handle)
{}

public:
//! @brief Constructs a \c async_memory_pool with the optionally specified initial pool size and release threshold.
//! @brief Constructs a \c device_memory_pool with the optionally specified initial pool size and release threshold.
//! If the pool size grows beyond the release threshold, unused memory held by the pool will be released at the next
//! synchronization event.
//! @throws cuda_error if the CUDA version does not support ``cudaMallocAsync``.
//! @param __device_id The device id of the device the stream pool is constructed on.
//! @param __pool_properties Optional, additional properties of the pool to be created.
explicit async_memory_pool(const ::cuda::experimental::device_ref __device_id,
async_memory_pool_properties __properties = {})
explicit device_memory_pool(const ::cuda::experimental::device_ref __device_id,
memory_pool_properties __properties = {})
: __pool_handle_(__create_cuda_mempool(__device_id.get(), __properties))
{}

//! @brief Disables construction from a plain `cudaMemPool_t`. We want to ensure clean ownership semantics.
async_memory_pool(::cudaMemPool_t) = delete;
device_memory_pool(::cudaMemPool_t) = delete;

async_memory_pool(async_memory_pool const&) = delete;
async_memory_pool(async_memory_pool&&) = delete;
async_memory_pool& operator=(async_memory_pool const&) = delete;
async_memory_pool& operator=(async_memory_pool&&) = delete;
device_memory_pool(device_memory_pool const&) = delete;
device_memory_pool(device_memory_pool&&) = delete;
device_memory_pool& operator=(device_memory_pool const&) = delete;
device_memory_pool& operator=(device_memory_pool&&) = delete;

//! @brief Destroys the \c async_memory_pool by releasing the internal ``cudaMemPool_t``.
~async_memory_pool() noexcept
//! @brief Destroys the \c device_memory_pool by releasing the internal ``cudaMemPool_t``.
~device_memory_pool() noexcept
{
_CCCL_ASSERT_CUDA_API(::cudaMemPoolDestroy, "~async_memory_pool() failed to destroy pool", __pool_handle_);
_CCCL_ASSERT_CUDA_API(::cudaMemPoolDestroy, "~device_memory_pool() failed to destroy pool", __pool_handle_);
}

//! @brief Tries to release memory.
Expand All @@ -263,7 +263,7 @@ public:
void trim_to(const size_t __min_bytes_to_keep)
{
_CCCL_TRY_CUDA_API(::cudaMemPoolTrimTo,
"Failed to call cudaMemPoolTrimTo in async_memory_pool::trim_to",
"Failed to call cudaMemPoolTrimTo in device_memory_pool::trim_to",
__pool_handle_,
__min_bytes_to_keep);
}
Expand All @@ -276,7 +276,7 @@ public:
size_t __value = 0;
_CCCL_TRY_CUDA_API(
::cudaMemPoolGetAttribute,
"Failed to call cudaMemPoolSetAttribute in async_memory_pool::get_attribute",
"Failed to call cudaMemPoolSetAttribute in device_memory_pool::get_attribute",
__pool_handle_,
__attr,
static_cast<void*>(&__value));
Expand All @@ -291,18 +291,18 @@ public:
{
if (__attr == ::cudaMemPoolAttrReservedMemCurrent || __attr == cudaMemPoolAttrUsedMemCurrent)
{
_CUDA_VSTD_NOVERSION::__throw_invalid_argument("Invalid attribute passed to async_memory_pool::set_attribute.");
_CUDA_VSTD_NOVERSION::__throw_invalid_argument("Invalid attribute passed to device_memory_pool::set_attribute.");
}
else if ((__attr == ::cudaMemPoolAttrReservedMemHigh || __attr == cudaMemPoolAttrUsedMemHigh) && __value != 0)
{
_CUDA_VSTD_NOVERSION::__throw_invalid_argument(
"async_memory_pool::set_attribute: It is illegal to set this "
"device_memory_pool::set_attribute: It is illegal to set this "
"attribute to a non-zero value.");
}

_CCCL_TRY_CUDA_API(
::cudaMemPoolSetAttribute,
"Failed to call cudaMemPoolSetAttribute in async_memory_pool::set_attribute",
"Failed to call cudaMemPoolSetAttribute in device_memory_pool::set_attribute",
__pool_handle_,
__attr,
static_cast<void*>(&__value));
Expand Down Expand Up @@ -355,17 +355,17 @@ public:
return ::cuda::experimental::mr::__mempool_get_access(__pool_handle_, __device);
}

//! @brief Equality comparison with another \c async_memory_pool.
//! @brief Equality comparison with another \c device_memory_pool.
//! @returns true if the stored ``cudaMemPool_t`` are equal.
_CCCL_NODISCARD constexpr bool operator==(async_memory_pool const& __rhs) const noexcept
_CCCL_NODISCARD constexpr bool operator==(device_memory_pool const& __rhs) const noexcept
{
return __pool_handle_ == __rhs.__pool_handle_;
}

# if _CCCL_STD_VER <= 2017
//! @brief Inequality comparison with another \c async_memory_pool.
//! @brief Inequality comparison with another \c device_memory_pool.
//! @returns true if the stored ``cudaMemPool_t`` are not equal.
_CCCL_NODISCARD constexpr bool operator!=(async_memory_pool const& __rhs) const noexcept
_CCCL_NODISCARD constexpr bool operator!=(device_memory_pool const& __rhs) const noexcept
{
return __pool_handle_ != __rhs.__pool_handle_;
}
Expand All @@ -374,26 +374,26 @@ public:
//! @brief Equality comparison with a \c cudaMemPool_t.
//! @param __rhs A \c cudaMemPool_t.
//! @returns true if the stored ``cudaMemPool_t`` is equal to \p __rhs.
_CCCL_NODISCARD_FRIEND constexpr bool operator==(async_memory_pool const& __lhs, ::cudaMemPool_t __rhs) noexcept
_CCCL_NODISCARD_FRIEND constexpr bool operator==(device_memory_pool const& __lhs, ::cudaMemPool_t __rhs) noexcept
{
return __lhs.__pool_handle_ == __rhs;
}

# if _CCCL_STD_VER <= 2017
//! @copydoc async_memory_pool::operator==(async_memory_pool const&, ::cudaMemPool_t)
_CCCL_NODISCARD_FRIEND constexpr bool operator==(::cudaMemPool_t __lhs, async_memory_pool const& __rhs) noexcept
//! @copydoc device_memory_pool::operator==(device_memory_pool const&, ::cudaMemPool_t)
_CCCL_NODISCARD_FRIEND constexpr bool operator==(::cudaMemPool_t __lhs, device_memory_pool const& __rhs) noexcept
{
return __rhs.__pool_handle_ == __lhs;
}

//! @copydoc async_memory_pool::operator==(async_memory_pool const&, ::cudaMemPool_t)
_CCCL_NODISCARD_FRIEND constexpr bool operator!=(async_memory_pool const& __lhs, ::cudaMemPool_t __rhs) noexcept
//! @copydoc device_memory_pool::operator==(device_memory_pool const&, ::cudaMemPool_t)
_CCCL_NODISCARD_FRIEND constexpr bool operator!=(device_memory_pool const& __lhs, ::cudaMemPool_t __rhs) noexcept
{
return __lhs.__pool_handle_ != __rhs;
}

//! @copydoc async_memory_pool::operator==(async_memory_pool const&, ::cudaMemPool_t)
_CCCL_NODISCARD_FRIEND constexpr bool operator!=(::cudaMemPool_t __lhs, async_memory_pool const& __rhs) noexcept
//! @copydoc device_memory_pool::operator==(device_memory_pool const&, ::cudaMemPool_t)
_CCCL_NODISCARD_FRIEND constexpr bool operator!=(::cudaMemPool_t __lhs, device_memory_pool const& __rhs) noexcept
{
return __rhs.__pool_handle_ != __lhs;
}
Expand All @@ -405,23 +405,23 @@ public:
return __pool_handle_;
}

//! @brief Construct an `async_memory_pool` object from a native `cudaMemPool_t` handle.
//! @brief Construct an `device_memory_pool` object from a native `cudaMemPool_t` handle.
//!
//! @param __handle The native handle
//!
//! @return The constructed `async_memory_pool` object
//! @return The constructed `device_memory_pool` object
//!
//! @note The constructed `async_memory_pool` object takes ownership of the native handle.
_CCCL_NODISCARD static async_memory_pool from_native_handle(::cudaMemPool_t __handle) noexcept
//! @note The constructed `device_memory_pool` object takes ownership of the native handle.
_CCCL_NODISCARD static device_memory_pool from_native_handle(::cudaMemPool_t __handle) noexcept
{
return async_memory_pool(__from_handle_t{}, __handle);
return device_memory_pool(__from_handle_t{}, __handle);
}

// Disallow construction from an `int`, e.g., `0`.
static async_memory_pool from_native_handle(int) = delete;
static device_memory_pool from_native_handle(int) = delete;

// Disallow construction from `nullptr`.
static async_memory_pool from_native_handle(_CUDA_VSTD::nullptr_t) = delete;
static device_memory_pool from_native_handle(_CUDA_VSTD::nullptr_t) = delete;
};

} // namespace cuda::experimental::mr
Expand All @@ -430,4 +430,4 @@ public:

#endif // !_CCCL_COMPILER_MSVC_2017 && !_CCCL_CUDACC_BELOW_11_2

#endif // _CUDAX__MEMORY_RESOURCE_CUDA_MEMORY_POOL
#endif // _CUDAX__MEMORY_RESOURCE_DEVICE_MEMORY_POOL
Loading

0 comments on commit 6edb860

Please sign in to comment.