Skip to content

Commit

Permalink
[FEA]: Use std::optional instead of thrust::optional (#1464)
Browse files Browse the repository at this point in the history
We already require c++17, so there is no benefit on pulling in `thrust::optional` as a dependency.

In addition thrust versiones its types with the cuda architecture, so using those types in API boundaries is not recommended.

This could be a potentially breaking change if a user would explicitly pass in `thrust::optional` into those APIs which seems unlikely

Authors:
  - Michael Schellenberger Costa (https://github.com/miscco)
  - Bradley Dice (https://github.com/bdice)

Approvers:
  - Mark Harris (https://github.com/harrism)
  - Bradley Dice (https://github.com/bdice)

URL: #1464
  • Loading branch information
miscco authored Feb 26, 2024
1 parent 51f0439 commit c408a32
Show file tree
Hide file tree
Showing 3 changed files with 174 additions and 30 deletions.
42 changes: 39 additions & 3 deletions include/rmm/mr/device/cuda_async_memory_resource.hpp
Original file line number Diff line number Diff line change
Expand Up @@ -23,11 +23,13 @@
#include <rmm/mr/device/cuda_async_view_memory_resource.hpp>
#include <rmm/mr/device/device_memory_resource.hpp>

#include <cuda/std/type_traits>
#include <cuda_runtime_api.h>
#include <thrust/optional.h>

#include <cstddef>
#include <limits>
#include <optional>

#if CUDART_VERSION >= 11020 // 11.2 introduced cudaMallocAsync
#ifndef RMM_DISABLE_CUDA_MALLOC_ASYNC
Expand Down Expand Up @@ -84,9 +86,43 @@ class cuda_async_memory_resource final : public device_memory_resource {
* `cudaMemHandleTypeNone` for no IPC support.
*/
// NOLINTNEXTLINE(bugprone-easily-swappable-parameters)
cuda_async_memory_resource(thrust::optional<std::size_t> initial_pool_size = {},
thrust::optional<std::size_t> release_threshold = {},
thrust::optional<allocation_handle_type> export_handle_type = {})
template <class Optional,
cuda::std::enable_if_t<cuda::std::is_same_v<cuda::std::remove_cvref_t<Optional>,
thrust::optional<std::size_t>>,
int> = 0>
[[deprecated("Use std::optional instead of thrust::optional.")]] //
explicit cuda_async_memory_resource(
Optional initial_pool_size,
Optional release_threshold = {},
thrust::optional<allocation_handle_type> export_handle_type = {})
: cuda_async_memory_resource(initial_pool_size.value_or(std::nullopt),
release_threshold.value_or(std::nullopt),
export_handle_type.value_or(std::nullopt))

{
}

/**
* @brief Constructs a cuda_async_memory_resource 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 rmm::logic_error if the CUDA version does not support `cudaMallocAsync`
*
* @param initial_pool_size Optional initial size in bytes of the pool. If no value is provided,
* initial pool size is half of the available GPU memory.
* @param release_threshold Optional release threshold size in bytes of the pool. If no value is
* provided, the release threshold is set to the total amount of memory on the current device.
* @param export_handle_type Optional `cudaMemAllocationHandleType` that allocations from this
* resource should support interprocess communication (IPC). Default is
* `cudaMemHandleTypeNone` for no IPC support.
*/
// NOLINTNEXTLINE(bugprone-easily-swappable-parameters)
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 = {})
{
#ifdef RMM_CUDA_MALLOC_ASYNC_SUPPORT
// Check if cudaMallocAsync Memory pool supported
Expand Down
142 changes: 131 additions & 11 deletions include/rmm/mr/device/pool_memory_resource.hpp
Original file line number Diff line number Diff line change
Expand Up @@ -25,6 +25,7 @@
#include <rmm/mr/device/detail/stream_ordered_memory_resource.hpp>
#include <rmm/mr/device/device_memory_resource.hpp>

#include <cuda/std/type_traits>
#include <cuda_runtime_api.h>
#include <thrust/iterator/counting_iterator.h>
#include <thrust/iterator/transform_iterator.h>
Expand All @@ -38,6 +39,7 @@
#include <map>
#include <mutex>
#include <numeric>
#include <optional>
#include <set>
#include <thread>
#include <unordered_map>
Expand Down Expand Up @@ -108,6 +110,37 @@ class pool_memory_resource final
friend class detail::stream_ordered_memory_resource<pool_memory_resource<Upstream>,
detail::coalescing_free_list>;

/**
* @brief Construct a `pool_memory_resource` and allocate the initial device memory
* pool using `upstream_mr`.
*
* @deprecated Use the constructor that takes an explicit initial pool size instead.
*
* @throws rmm::logic_error if `upstream_mr == nullptr`
* @throws rmm::logic_error if `initial_pool_size` is neither the default nor aligned to a
* multiple of pool_memory_resource::allocation_alignment bytes.
* @throws rmm::logic_error if `maximum_pool_size` is neither the default nor aligned to a
* multiple of pool_memory_resource::allocation_alignment bytes.
*
* @param upstream_mr The memory_resource from which to allocate blocks for the pool.
* @param initial_pool_size Minimum size, in bytes, of the initial pool. Defaults to zero.
* @param maximum_pool_size Maximum size, in bytes, that the pool can grow to. Defaults to all
* of the available memory from the upstream resource.
*/
template <class Optional,
cuda::std::enable_if_t<cuda::std::is_same_v<cuda::std::remove_cvref_t<Optional>,
thrust::optional<std::size_t>>,
int> = 0>
[[deprecated(
"Must specify initial_pool_size and use std::optional instead of thrust::optional.")]] //
explicit pool_memory_resource(Upstream* upstream_mr,
Optional initial_pool_size,
Optional maximum_pool_size = thrust::nullopt)
: pool_memory_resource(
upstream_mr, initial_pool_size.value_or(0), maximum_pool_size.value_or(std::nullopt))
{
}

/**
* @brief Construct a `pool_memory_resource` and allocate the initial device memory
* pool using `upstream_mr`.
Expand All @@ -127,12 +160,43 @@ class pool_memory_resource final
*/
[[deprecated("Must specify initial_pool_size")]] //
explicit pool_memory_resource(Upstream* upstream_mr,
thrust::optional<std::size_t> initial_pool_size = thrust::nullopt,
thrust::optional<std::size_t> maximum_pool_size = thrust::nullopt)
std::optional<std::size_t> initial_pool_size = std::nullopt,
std::optional<std::size_t> maximum_pool_size = std::nullopt)
: pool_memory_resource(upstream_mr, initial_pool_size.value_or(0), maximum_pool_size)
{
}

/**
* @brief Construct a `pool_memory_resource` and allocate the initial device memory pool using
* `upstream_mr`.
*
* @deprecated Use the constructor that takes an explicit initial pool size instead.
*
* @throws rmm::logic_error if `upstream_mr == nullptr`
* @throws rmm::logic_error if `initial_pool_size` is neither the default nor aligned to a
* multiple of pool_memory_resource::allocation_alignment bytes.
* @throws rmm::logic_error if `maximum_pool_size` is neither the default nor aligned to a
* multiple of pool_memory_resource::allocation_alignment bytes.
*
* @param upstream_mr The memory_resource from which to allocate blocks for the pool.
* @param initial_pool_size Minimum size, in bytes, of the initial pool. Defaults to zero.
* @param maximum_pool_size Maximum size, in bytes, that the pool can grow to. Defaults to all
* of the available memory from the upstream resource.
*/
template <class Optional,
cuda::std::enable_if_t<cuda::std::is_same_v<cuda::std::remove_cvref_t<Optional>,
thrust::optional<std::size_t>>,
int> = 0>
[[deprecated(
"Must specify initial_pool_size and use std::optional instead of thrust::optional.")]] //
explicit pool_memory_resource(Upstream& upstream_mr,
Optional initial_pool_size,
Optional maximum_pool_size = thrust::nullopt)
: pool_memory_resource(
upstream_mr, initial_pool_size.value_or(0), maximum_pool_size.value_or(std::nullopt))
{
}

/**
* @brief Construct a `pool_memory_resource` and allocate the initial device memory pool using
* `upstream_mr`.
Expand All @@ -154,12 +218,39 @@ class pool_memory_resource final
cuda::std::enable_if_t<cuda::mr::async_resource<Upstream2>, int> = 0>
[[deprecated("Must specify initial_pool_size")]] //
explicit pool_memory_resource(Upstream2& upstream_mr,
thrust::optional<std::size_t> initial_pool_size = thrust::nullopt,
thrust::optional<std::size_t> maximum_pool_size = thrust::nullopt)
std::optional<std::size_t> initial_pool_size = std::nullopt,
std::optional<std::size_t> maximum_pool_size = std::nullopt)
: pool_memory_resource(upstream_mr, initial_pool_size.value_or(0), maximum_pool_size)
{
}

/**
* @brief Construct a `pool_memory_resource` and allocate the initial device memory pool using
* `upstream_mr`.
*
* @throws rmm::logic_error if `upstream_mr == nullptr`
* @throws rmm::logic_error if `initial_pool_size` is not aligned to a multiple of
* pool_memory_resource::allocation_alignment bytes.
* @throws rmm::logic_error if `maximum_pool_size` is neither the default nor aligned to a
* multiple of pool_memory_resource::allocation_alignment bytes.
*
* @param upstream_mr The memory_resource from which to allocate blocks for the pool.
* @param initial_pool_size Minimum size, in bytes, of the initial pool.
* @param maximum_pool_size Maximum size, in bytes, that the pool can grow to. Defaults to all
* of the available from the upstream resource.
*/
template <class Optional,
cuda::std::enable_if_t<cuda::std::is_same_v<cuda::std::remove_cvref_t<Optional>,
thrust::optional<std::size_t>>,
int> = 0>
[[deprecated("Use std::optional instead of thrust::optional.")]] //
explicit pool_memory_resource(Upstream* upstream_mr,
std::size_t initial_pool_size,
Optional maximum_pool_size)
: pool_memory_resource(upstream_mr, initial_pool_size, maximum_pool_size.value_or(std::nullopt))
{
}

/**
* @brief Construct a `pool_memory_resource` and allocate the initial device memory pool using
* `upstream_mr`.
Expand All @@ -177,7 +268,7 @@ class pool_memory_resource final
*/
explicit pool_memory_resource(Upstream* upstream_mr,
std::size_t initial_pool_size,
thrust::optional<std::size_t> maximum_pool_size = thrust::nullopt)
std::optional<std::size_t> maximum_pool_size = std::nullopt)
: upstream_mr_{[upstream_mr]() {
RMM_EXPECTS(nullptr != upstream_mr, "Unexpected null upstream pointer.");
return upstream_mr;
Expand All @@ -191,6 +282,35 @@ class pool_memory_resource final
initialize_pool(initial_pool_size, maximum_pool_size);
}

/**
* @brief Construct a `pool_memory_resource` and allocate the initial device memory pool using
* `upstream_mr`.
*
* @throws rmm::logic_error if `upstream_mr == nullptr`
* @throws rmm::logic_error if `initial_pool_size` is not aligned to a multiple of
* pool_memory_resource::allocation_alignment bytes.
* @throws rmm::logic_error if `maximum_pool_size` is neither the default nor aligned to a
* multiple of pool_memory_resource::allocation_alignment bytes.
*
* @param upstream_mr The memory_resource from which to allocate blocks for the pool.
* @param initial_pool_size Minimum size, in bytes, of the initial pool.
* @param maximum_pool_size Maximum size, in bytes, that the pool can grow to. Defaults to all
* of the available memory from the upstream resource.
*/
template <class Optional,
cuda::std::enable_if_t<cuda::std::is_same_v<cuda::std::remove_cvref_t<Optional>,
thrust::optional<std::size_t>>,
int> = 0>
[[deprecated("Use std::optional instead of thrust::optional.")]] //
explicit pool_memory_resource(Upstream& upstream_mr,
std::size_t initial_pool_size,
Optional maximum_pool_size)
: pool_memory_resource(cuda::std::addressof(upstream_mr),
initial_pool_size,
maximum_pool_size.value_or(std::nullopt))
{
}

/**
* @brief Construct a `pool_memory_resource` and allocate the initial device memory pool using
* `upstream_mr`.
Expand All @@ -210,7 +330,7 @@ class pool_memory_resource final
cuda::std::enable_if_t<cuda::mr::async_resource<Upstream2>, int> = 0>
explicit pool_memory_resource(Upstream2& upstream_mr,
std::size_t initial_pool_size,
thrust::optional<std::size_t> maximum_pool_size = thrust::nullopt)
std::optional<std::size_t> maximum_pool_size = std::nullopt)
: pool_memory_resource(cuda::std::addressof(upstream_mr), initial_pool_size, maximum_pool_size)
{
}
Expand Down Expand Up @@ -312,7 +432,7 @@ class pool_memory_resource final
*
* @throws logic_error if @p initial_size is larger than @p maximum_size (if set).
*/
void initialize_pool(std::size_t initial_size, thrust::optional<std::size_t> maximum_size)
void initialize_pool(std::size_t initial_size, std::optional<std::size_t> maximum_size)
{
current_pool_size_ = 0; // try_to_expand will set this if it succeeds
maximum_pool_size_ = maximum_size;
Expand Down Expand Up @@ -376,18 +496,18 @@ class pool_memory_resource final
* @param stream The stream on which the memory is to be used.
* @return block_type The allocated block
*/
thrust::optional<block_type> block_from_upstream(std::size_t size, cuda_stream_view stream)
std::optional<block_type> block_from_upstream(std::size_t size, cuda_stream_view stream)
{
RMM_LOG_DEBUG("[A][Stream {}][Upstream {}B]", fmt::ptr(stream.value()), size);

if (size == 0) { return {}; }

try {
void* ptr = get_upstream()->allocate_async(size, stream);
return thrust::optional<block_type>{
return std::optional<block_type>{
*upstream_blocks_.emplace(static_cast<char*>(ptr), size, true).first};
} catch (std::exception const& e) {
return thrust::nullopt;
return std::nullopt;
}
}

Expand Down Expand Up @@ -515,7 +635,7 @@ class pool_memory_resource final
private:
Upstream* upstream_mr_; // The "heap" to allocate the pool from
std::size_t current_pool_size_{};
thrust::optional<std::size_t> maximum_pool_size_{};
std::optional<std::size_t> maximum_pool_size_{};

#ifdef RMM_POOL_TRACK_ALLOCATIONS
std::set<block_type, rmm::mr::detail::compare_blocks<block_type>> allocated_blocks_;
Expand Down
20 changes: 4 additions & 16 deletions python/rmm/_lib/memory_resource.pyx
Original file line number Diff line number Diff line change
Expand Up @@ -25,6 +25,7 @@ from libc.stddef cimport size_t
from libc.stdint cimport int8_t, int64_t, uintptr_t
from libcpp cimport bool
from libcpp.memory cimport make_unique, unique_ptr
from libcpp.optional cimport optional
from libcpp.pair cimport pair
from libcpp.string cimport string

Expand Down Expand Up @@ -76,19 +77,6 @@ cdef extern from *:

# NOTE: Keep extern declarations in .pyx file as much as possible to avoid
# leaking dependencies when importing RMM Cython .pxd files
cdef extern from "thrust/optional.h" namespace "thrust" nogil:

struct nullopt_t:
pass

cdef nullopt_t nullopt

cdef cppclass optional[T]:
optional()
optional(T v)

cdef optional[T] make_optional[T](T v)

cdef extern from "rmm/mr/device/cuda_memory_resource.hpp" \
namespace "rmm::mr" nogil:
cdef cppclass cuda_memory_resource(device_memory_resource):
Expand Down Expand Up @@ -322,13 +310,13 @@ cdef class CudaAsyncMemoryResource(DeviceMemoryResource):
cdef optional[size_t] c_initial_pool_size = (
optional[size_t]()
if initial_pool_size is None
else optional[size_t](initial_pool_size)
else optional[size_t](<size_t> initial_pool_size)
)

cdef optional[size_t] c_release_threshold = (
optional[size_t]()
if release_threshold is None
else optional[size_t](release_threshold)
else optional[size_t](<size_t> release_threshold)
)

# If IPC memory handles are not supported, the constructor below will
Expand Down Expand Up @@ -382,7 +370,7 @@ cdef class PoolMemoryResource(UpstreamResourceAdaptor):
c_maximum_pool_size = (
optional[size_t]() if
maximum_pool_size is None
else make_optional[size_t](maximum_pool_size)
else optional[size_t](<size_t> maximum_pool_size)
)
self.c_obj.reset(
new pool_memory_resource[device_memory_resource](
Expand Down

0 comments on commit c408a32

Please sign in to comment.