From a5d8df2e1b2e573e9bd37b155b6213f3b88732a0 Mon Sep 17 00:00:00 2001 From: niranda perera Date: Wed, 29 Jul 2026 17:52:57 -0700 Subject: [PATCH 1/6] porting impl from rapidsmpf Signed-off-by: niranda perera --- .../reservation_aware_resource_adaptor.hpp | 638 ++++++++++++++++++ src/memory/CMakeLists.txt | 32 +- ...tal_reservation_aware_resource_adaptor.cpp | 85 +++ test/CMakeLists.txt | 1 + ...tal_reservation_aware_resource_adaptor.cpp | 292 ++++++++ 5 files changed, 1033 insertions(+), 15 deletions(-) create mode 100644 include/cucascade/memory/experimental/reservation_aware_resource_adaptor.hpp create mode 100644 src/memory/experimental_reservation_aware_resource_adaptor.cpp create mode 100644 test/memory/test_experimental_reservation_aware_resource_adaptor.cpp diff --git a/include/cucascade/memory/experimental/reservation_aware_resource_adaptor.hpp b/include/cucascade/memory/experimental/reservation_aware_resource_adaptor.hpp new file mode 100644 index 0000000..3957d31 --- /dev/null +++ b/include/cucascade/memory/experimental/reservation_aware_resource_adaptor.hpp @@ -0,0 +1,638 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. + * All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#pragma once + +#include +#include + +#include +#include +#include +#include + +#include +#include + +#include +#include +#include +#include +#include +#include +#include +#include + +namespace cucascade { +namespace memory { +namespace experimental { + +class reservation_aware_resource_adaptor; +class memory_reservation; + +using any_device_resource = ::cuda::mr::any_resource<::cuda::mr::device_accessible>; + +/** + * @brief Snapshot of the adaptor's main (non-scoped) allocation accounting. + */ +struct memory_record { + std::int64_t num_current_allocs{0}; + std::int64_t num_total_allocs{0}; + std::int64_t current{0}; + std::int64_t total{0}; + std::int64_t peak{0}; + std::int64_t max{0}; +}; + +/** + * @brief Policy controlling whether a reservation may exceed the adaptor's limit. + */ +enum class allow_overbooking : bool { + NO, ///< Fail the request rather than exceed the limit. + YES, ///< Grant the request even when the memory isn't available. +}; + +namespace detail { + +template +[[nodiscard]] constexpr To safe_cast(From value) +{ + if constexpr (std::is_same_v) { + return value; + } else { + if (!std::in_range(value)) { + throw std::overflow_error("cucascade cast: value out of range " + std::to_string(value)); + } + return static_cast(value); + } +} + +// The adaptor is a template parameter only so that the reservation can hold one by +// value: `reservation_aware_resource_adaptor` is defined in terms of the adaptor impl +// below, so it is still incomplete here, and only a dependent member type defers the +// completeness requirement to instantiation time. +template + requires std::same_as +class memory_reservation_impl; + +/** + * @brief Shared state of a reservation-aware resource adaptor. + * + * Owns an upstream device resource and tracks a single main memory record (no scoped + * records). Reservations are granted against a runtime-adjustable limit. + */ +class reservation_aware_resource_adaptor_impl { + public: + /** + * @brief Construct with a primary memory resource and a memory limit. + * + * @param upstream_mr The primary memory resource (moved in). + * @param limit Maximum number of bytes that may be allocated and reserved. + */ + reservation_aware_resource_adaptor_impl(any_device_resource upstream_mr, std::int64_t limit) + : upstream_mr_{std::move(upstream_mr)}, limit_{limit} + { + } + + ~reservation_aware_resource_adaptor_impl() = default; + + reservation_aware_resource_adaptor_impl(reservation_aware_resource_adaptor_impl const&) = delete; + reservation_aware_resource_adaptor_impl(reservation_aware_resource_adaptor_impl&&) = delete; + reservation_aware_resource_adaptor_impl& operator=( + reservation_aware_resource_adaptor_impl const&) = delete; + reservation_aware_resource_adaptor_impl& operator=(reservation_aware_resource_adaptor_impl&&) = + delete; + + [[nodiscard]] bool operator==(reservation_aware_resource_adaptor_impl const& other) const noexcept + { + return this == std::addressof(other); + } + + [[nodiscard]] any_device_resource const& get_upstream_resource() const noexcept + { + return upstream_mr_; + } + + [[nodiscard]] std::int64_t limit() const noexcept + { + return limit_.load(std::memory_order_acquire); + } + + void set_limit(std::int64_t limit) noexcept { limit_.store(limit, std::memory_order_release); } + + [[nodiscard]] std::int64_t total_reserved() const noexcept + { + return total_reserved_.load(std::memory_order_acquire); + } + + [[nodiscard]] std::int64_t current_allocated() const noexcept + { + return current_.load(std::memory_order_acquire); + } + + [[nodiscard]] std::int64_t available() const noexcept + { + return limit() - current_allocated() - total_reserved(); + } + + [[nodiscard]] memory_record get_main_record() const + { + return memory_record{ + .num_current_allocs = num_current_allocs_.load(std::memory_order_acquire), + .num_total_allocs = num_total_allocs_.load(std::memory_order_acquire), + .current = current_.load(std::memory_order_acquire), + .total = total_.load(std::memory_order_acquire), + .peak = peak_.peak(), + .max = max_.peak(), + }; + } + + /** + * @brief Reserve @p size bytes against the limit. + * + * @param size The number of bytes to reserve. + * @param allow_overbooking Whether to grant the reservation even when the memory + * isn't available. + * @return A pair of the number of bytes granted (either @p size or zero) and the + * number of bytes by which the request overbooks the limit. + * + * @note Rejections are best-effort under contention: concurrent requests each claim + * before checking, so requests that would fit individually may be rejected. + */ + [[nodiscard]] std::pair reserve(std::size_t size, + bool allow_overbooking) + { + auto const want = safe_cast(size); + std::int64_t const capacity = limit() - current_allocated(); + + // Claim the bytes up front and roll back if they didn't fit. While a claim is + // being rolled back the reserved total reads high, which makes a concurrent + // `available()` pessimistic, never optimistic. + auto const reserved = total_reserved_.add(want, std::memory_order_acq_rel) - want; + std::int64_t const headroom = capacity - (reserved + want); + if (headroom >= 0) { return {size, 0}; } + auto const overbooking = safe_cast(-headroom); + if (!allow_overbooking) { + total_reserved_.sub(want, std::memory_order_acq_rel); + return {0, overbooking}; + } + return {size, overbooking}; + } + + void* allocate(::cuda::stream_ref stream, + std::size_t bytes, + std::size_t alignment = rmm::CUDA_ALLOCATION_ALIGNMENT) + { + void* ret = upstream_mr_.allocate(stream, bytes, alignment); + record_allocation(safe_cast(bytes)); + return ret; + } + + void deallocate(::cuda::stream_ref stream, + void* ptr, + std::size_t bytes, + std::size_t alignment = rmm::CUDA_ALLOCATION_ALIGNMENT) noexcept + { + record_deallocation(safe_cast(bytes)); + upstream_mr_.deallocate(stream, ptr, bytes, alignment); + } + + void* allocate_sync(std::size_t bytes, std::size_t alignment = rmm::CUDA_ALLOCATION_ALIGNMENT) + { + auto* ptr = allocate(sync_stream_, bytes, alignment); + sync_stream_.synchronize(); + return ptr; + } + + void deallocate_sync(void* ptr, + std::size_t bytes, + std::size_t alignment = rmm::CUDA_ALLOCATION_ALIGNMENT) noexcept + { + deallocate(sync_stream_, ptr, bytes, alignment); + sync_stream_.synchronize_no_throw(); + } + + friend void get_property(reservation_aware_resource_adaptor_impl const&, + ::cuda::mr::device_accessible) noexcept + { + } + + private: + friend class memory_reservation_impl; + + void record_allocation(std::int64_t nbytes) + { + num_total_allocs_.add(1, std::memory_order_acq_rel); + num_current_allocs_.add(1, std::memory_order_acq_rel); + auto const current = current_.add(nbytes, std::memory_order_acq_rel); + total_.add(nbytes, std::memory_order_acq_rel); + peak_.update_peak(current); + max_.update_peak(nbytes); + } + + void record_deallocation(std::int64_t nbytes) noexcept + { + current_.sub(nbytes, std::memory_order_acq_rel); + num_current_allocs_.sub(1, std::memory_order_acq_rel); + } + + any_device_resource upstream_mr_; + + utils::atomic_bounded_counter num_current_allocs_{0}; + utils::atomic_bounded_counter num_total_allocs_{0}; + utils::atomic_bounded_counter current_{0}; + utils::atomic_bounded_counter total_{0}; + utils::atomic_peak_tracker peak_; + utils::atomic_peak_tracker max_; ///< Largest single allocation observed. + + std::atomic limit_; + // Reservations move bytes in and out of this counter as they allocate, free, and die. + utils::atomic_bounded_counter total_reserved_{0}; + + rmm::cuda_stream sync_stream_{rmm::cuda_stream::flags::non_blocking}; +}; + +/** + * @brief Shared state of a memory reservation. + * + * Satisfies the `cuda::mr::resource` concept so it can be a `cuda::mr::shared_resource`. + * Allocating moves bytes from the adaptor's reserved counter to its allocated counter; + * the unspent balance is refunded only when the last reference dies. + * + * @tparam Adaptor Always `reservation_aware_resource_adaptor`. + */ +template + requires std::same_as +class memory_reservation_impl { + public: + memory_reservation_impl(Adaptor adaptor, std::int64_t grant, std::size_t overbooking) + : adaptor_{std::move(adaptor)}, grant_{grant}, overbooking_{overbooking}, balance_{grant} + { + } + + ~memory_reservation_impl() + { + adaptor_->total_reserved_.sub(balance(), std::memory_order_acq_rel); + } + + memory_reservation_impl(memory_reservation_impl const&) = delete; + memory_reservation_impl(memory_reservation_impl&&) = delete; + memory_reservation_impl& operator=(memory_reservation_impl const&) = delete; + memory_reservation_impl& operator=(memory_reservation_impl&&) = delete; + + [[nodiscard]] std::int64_t grant() const noexcept { return grant_; } + + [[nodiscard]] std::size_t overbooking() const noexcept { return overbooking_; } + + [[nodiscard]] std::int64_t balance() const noexcept + { + return balance_.load(std::memory_order_acquire); + } + + void* allocate(::cuda::stream_ref stream, + std::size_t bytes, + std::size_t alignment = rmm::CUDA_ALLOCATION_ALIGNMENT) + { + auto const amount = safe_cast(bytes); + draw_down_res(amount); + void* ptr = nullptr; + try { + ptr = adaptor_->allocate(stream, bytes, alignment); + } catch (...) { + balance_.fetch_add(amount, std::memory_order_acq_rel); + throw; + } + adaptor_->total_reserved_.sub(amount, std::memory_order_acq_rel); + return ptr; + } + + void deallocate(::cuda::stream_ref stream, + void* ptr, + std::size_t bytes, + std::size_t alignment = rmm::CUDA_ALLOCATION_ALIGNMENT) noexcept + { + auto const amount = safe_cast(bytes); + balance_.fetch_add(amount, std::memory_order_acq_rel); + adaptor_->total_reserved_.add(amount, std::memory_order_acq_rel); + adaptor_->deallocate(stream, ptr, bytes, alignment); + } + + void* allocate_sync(std::size_t bytes, std::size_t alignment = rmm::CUDA_ALLOCATION_ALIGNMENT) + { + auto* ptr = allocate(adaptor_->sync_stream_, bytes, alignment); + adaptor_->sync_stream_.synchronize(); + return ptr; + } + + void deallocate_sync(void* ptr, + std::size_t bytes, + std::size_t alignment = rmm::CUDA_ALLOCATION_ALIGNMENT) noexcept + { + deallocate(adaptor_->sync_stream_, ptr, bytes, alignment); + adaptor_->sync_stream_.synchronize_no_throw(); + } + + [[nodiscard]] bool operator==(memory_reservation_impl const& other) const noexcept + { + return this == std::addressof(other); + } + + friend void get_property(memory_reservation_impl const&, ::cuda::mr::device_accessible) noexcept + { + } + + [[nodiscard]] Adaptor const& adaptor() const noexcept { return adaptor_; } + + private: + void draw_down_res(std::int64_t bytes) + { + auto balance = balance_.load(std::memory_order_relaxed); + do { + if (bytes > balance) { + CUCASCADE_FAIL("allocation of " + std::to_string(bytes) + + " bytes exceeds reservation (grant: " + std::to_string(grant_) + + ", remaining: " + std::to_string(balance) + ")", + rmm::out_of_memory); + } + } while (!balance_.compare_exchange_weak( + balance, balance - bytes, std::memory_order_acq_rel, std::memory_order_relaxed)); + } + + Adaptor adaptor_; + std::int64_t const grant_; + std::size_t const overbooking_; + std::atomic balance_; +}; + +} // namespace detail + +/** + * @brief A memory resource adaptor that only allocates through reservations. + * + * This adaptor wraps a primary device memory resource and adds a memory limit with + * allocation tracking. Memory is obtained by calling `reserve()` and allocating through + * the returned `memory_reservation`. + * + * This class is copyable and shares ownership of its internal state via + * `cuda::mr::shared_resource`. + * + * @par Allocating without a reservation + * + * The adaptor is itself a memory resource, so it can be handed to cudf or an RMM + * container directly. Those allocations are tracked, and therefore still consume + * `available()`, but they draw on no reservation and are not capped. Allocate through + * a `memory_reservation` when the budget has to be enforced. + * + * @par Accounting + * + * Three quantities describe the state of the adaptor: + * - `current_allocated()`: bytes currently allocated, tracked on every allocation. + * - `total_reserved()`: bytes held by live reservations but not yet allocated. + * - `available() == limit() - current_allocated() - total_reserved()`. + * + * Allocating through a reservation moves bytes from the second bucket to the first, + * leaving `available()` unchanged; that is what makes a reservation a promise. + */ +class reservation_aware_resource_adaptor + : public ::cuda::mr::shared_resource { + public: + /// @brief The adaptor's shared implementation. + using impl_type = detail::reservation_aware_resource_adaptor_impl; + + /// @brief The reference-counted handle on the shared implementation. + using shared_base = ::cuda::mr::shared_resource; + + /// @brief Tag this resource as device-accessible for the CCCL concept. + friend void get_property(reservation_aware_resource_adaptor const&, + ::cuda::mr::device_accessible) noexcept + { + } + + /** + * @brief Construct with the specified primary memory resource and limit. + * + * @param upstream_mr The primary memory resource. + * @param limit Maximum number of bytes that may be allocated and reserved. + */ + reservation_aware_resource_adaptor(any_device_resource upstream_mr, std::int64_t limit); + + /** + * @brief Equality comparison. + * + * @param other The other adaptor to compare. + * @return True if both adaptors share the same underlying state. + */ + [[nodiscard]] bool operator==(reservation_aware_resource_adaptor const& other) const noexcept + { + return get() == other.get(); + } + + /** + * @brief Reserve an amount of memory. + * + * Creates a new reservation of the specified size to inform about upcoming + * allocations. + * + * If overbooking is allowed, a reservation of @p size is returned even when the + * memory isn't available. In that case the caller must free (at least) + * `memory_reservation::overbooking()` bytes before using the reservation. + * + * If overbooking isn't allowed, a reservation of size zero is returned on failure, + * with `memory_reservation::overbooking()` reporting by how much the request missed. + * A zero-sized reservation fails at allocation time: the first allocation through it + * throws `rmm::out_of_memory`. + * + * @param size The number of bytes to reserve. + * @param overbooking_policy Whether overbooking is allowed. + * @return The reservation. On success its grant always equals @p size and on + * failure it always equals zero (a zero-sized reservation never fails). + */ + [[nodiscard]] memory_reservation reserve(std::size_t size, allow_overbooking overbooking_policy); + + /** + * @brief Get the memory limit. + * + * @return The limit in bytes. + */ + [[nodiscard]] std::int64_t limit() const noexcept; + + /** + * @brief Update the memory limit at runtime. + * + * @param limit The new byte limit. + */ + void set_limit(std::int64_t limit) noexcept; + + /** + * @brief Get the total current allocated memory through this adaptor. + * + * @return Total number of currently allocated bytes. + */ + [[nodiscard]] std::int64_t current_allocated() const noexcept; + + /** + * @brief Get the memory currently held by live reservations. + * + * Excludes reserved bytes that have already been allocated; those are reported by + * `current_allocated()` instead. + * + * @return Total number of reserved bytes. + */ + [[nodiscard]] std::int64_t total_reserved() const noexcept; + + /** + * @brief Get the memory available for new reservations. + * + * Computed as `limit() - current_allocated() - total_reserved()`. May be negative + * when reservations have overbooked the limit. + * + * @return The available memory in bytes. + */ + [[nodiscard]] std::int64_t available() const noexcept; + + /** + * @brief Returns a snapshot of the main memory record. + * + * @return A copy of the current main memory record. + */ + [[nodiscard]] memory_record get_main_record() const; + + /** + * @brief Get a reference to the primary upstream resource. + * + * @return Reference to the RMM memory resource. + */ + [[nodiscard]] rmm::device_async_resource_ref get_upstream_resource() const noexcept; +}; + +static_assert( + ::cuda::mr::resource_with); + +/** + * @brief A memory reservation that is itself a memory resource. + * + * Granted by `reservation_aware_resource_adaptor::reserve()`, a reservation holds a + * budget of bytes carved out of the adaptor's limit. It is an RMM memory resource, so + * it can be handed to cudf (or anything else taking a `rmm::device_async_resource_ref`), + * and every allocation made through it is charged against that budget. An allocation + * exceeding the remaining `balance()` throws `rmm::out_of_memory`; deallocating returns + * the bytes to the balance. + * + * @par Ownership + * + * Like `reservation_aware_resource_adaptor`, this is a `cuda::mr::shared_resource`, so + * copies share the same reservation and are interchangeable. RMM stores such a copy + * inside every buffer allocated from the reservation, which is what keeps the + * reservation alive for as long as those buffers need it to service deallocations. + * + * The unspent balance is refunded to the adaptor when the last copy dies. Reserving + * more than is allocated therefore keeps the surplus out of circulation for as long as + * any derived buffer lives, so reserve what you actually use. + * + * @code{.cpp} + * auto res = adaptor.reserve(1 << 30, allow_overbooking::NO); + * auto table = cudf::groupby(..., stream, res); + * @endcode + */ +class memory_reservation : public ::cuda::mr::shared_resource< + detail::memory_reservation_impl> { + using shared_base = ::cuda::mr::shared_resource< + detail::memory_reservation_impl>; + + public: + /// @brief The shared state of the reservation. + using impl_type = detail::memory_reservation_impl; + + /// @brief Tag this resource as device-accessible for the CCCL concept. + friend void get_property(memory_reservation const&, ::cuda::mr::device_accessible) noexcept {} + + /** + * @brief Equality comparison. + * + * @param other The other reservation to compare. + * @return True if both refer to the same reservation. + */ + [[nodiscard]] bool operator==(memory_reservation const& other) const noexcept + { + return get() == other.get(); + } + + /** + * @brief The number of bytes originally granted. + * + * @return The granted size in bytes. + */ + [[nodiscard]] std::size_t grant() const noexcept + { + return detail::safe_cast(get().grant()); + } + + /** + * @brief The remaining unallocated size of the reservation. + * + * @return The remaining size in bytes. + */ + [[nodiscard]] std::size_t balance() const noexcept + { + return detail::safe_cast(get().balance()); + } + + /** + * @brief The number of bytes by which the grant overbooks the adaptor's limit. + * + * Nonzero only when the reservation was granted with `allow_overbooking::YES`. The + * caller must free at least this much memory before using the reservation. + * + * @return The overbooked size in bytes. + */ + [[nodiscard]] std::size_t overbooking() const noexcept { return get().overbooking(); } + + /** + * @brief The adaptor that granted the reservation. + * + * @return The adaptor. + */ + [[nodiscard]] reservation_aware_resource_adaptor const& adaptor() const noexcept + { + return get().adaptor(); + } + + private: + friend class reservation_aware_resource_adaptor; + + /** + * @brief Construct from an already-granted reservation. + * + * Private so that only `reservation_aware_resource_adaptor` can grant reservations. + * The reservation holds a copy of the adaptor, so the adaptor stays alive for as + * long as any buffer allocated from the reservation needs it. + * + * @param adaptor The adaptor that granted the reservation. + * @param granted The number of bytes granted. + * @param overbooking The number of bytes by which @p granted overbooks the limit. + */ + memory_reservation(reservation_aware_resource_adaptor const& adaptor, + std::size_t granted, + std::size_t overbooking); +}; + +static_assert(::cuda::mr::resource_with); + +} // namespace experimental +} // namespace memory +} // namespace cucascade diff --git a/src/memory/CMakeLists.txt b/src/memory/CMakeLists.txt index 03ba85e..007c2f8 100644 --- a/src/memory/CMakeLists.txt +++ b/src/memory/CMakeLists.txt @@ -18,21 +18,23 @@ if(TARGET cucascade_objects) target_sources( cucascade_objects - PRIVATE ${CMAKE_CURRENT_SOURCE_DIR}/common.cpp - ${CMAKE_CURRENT_SOURCE_DIR}/disk_access_limiter.cpp - ${CMAKE_CURRENT_SOURCE_DIR}/fixed_size_host_memory_resource.cpp - ${CMAKE_CURRENT_SOURCE_DIR}/memory_reservation.cpp - ${CMAKE_CURRENT_SOURCE_DIR}/memory_reservation_manager.cpp - ${CMAKE_CURRENT_SOURCE_DIR}/memory_space.cpp - ${CMAKE_CURRENT_SOURCE_DIR}/notification_channel.cpp - ${CMAKE_CURRENT_SOURCE_DIR}/null_device_memory_resource.cpp - ${CMAKE_CURRENT_SOURCE_DIR}/numa_region_pinned_host_allocator.cpp - ${CMAKE_CURRENT_SOURCE_DIR}/oom_handling_policy.cpp - ${CMAKE_CURRENT_SOURCE_DIR}/reservation_aware_resource_adaptor.cpp - ${CMAKE_CURRENT_SOURCE_DIR}/reservation_manager_configurator.cpp - ${CMAKE_CURRENT_SOURCE_DIR}/small_pinned_host_memory_resource.cpp - ${CMAKE_CURRENT_SOURCE_DIR}/stream_pool.cpp - ${CMAKE_CURRENT_SOURCE_DIR}/error.cpp) + PRIVATE + ${CMAKE_CURRENT_SOURCE_DIR}/common.cpp + ${CMAKE_CURRENT_SOURCE_DIR}/disk_access_limiter.cpp + ${CMAKE_CURRENT_SOURCE_DIR}/fixed_size_host_memory_resource.cpp + ${CMAKE_CURRENT_SOURCE_DIR}/memory_reservation.cpp + ${CMAKE_CURRENT_SOURCE_DIR}/memory_reservation_manager.cpp + ${CMAKE_CURRENT_SOURCE_DIR}/memory_space.cpp + ${CMAKE_CURRENT_SOURCE_DIR}/notification_channel.cpp + ${CMAKE_CURRENT_SOURCE_DIR}/null_device_memory_resource.cpp + ${CMAKE_CURRENT_SOURCE_DIR}/numa_region_pinned_host_allocator.cpp + ${CMAKE_CURRENT_SOURCE_DIR}/oom_handling_policy.cpp + ${CMAKE_CURRENT_SOURCE_DIR}/reservation_aware_resource_adaptor.cpp + ${CMAKE_CURRENT_SOURCE_DIR}/experimental_reservation_aware_resource_adaptor.cpp + ${CMAKE_CURRENT_SOURCE_DIR}/reservation_manager_configurator.cpp + ${CMAKE_CURRENT_SOURCE_DIR}/small_pinned_host_memory_resource.cpp + ${CMAKE_CURRENT_SOURCE_DIR}/stream_pool.cpp + ${CMAKE_CURRENT_SOURCE_DIR}/error.cpp) endif() target_sources(cucascade_topology_discovery_objects PRIVATE ${CMAKE_CURRENT_SOURCE_DIR}/topology_discovery.cpp) diff --git a/src/memory/experimental_reservation_aware_resource_adaptor.cpp b/src/memory/experimental_reservation_aware_resource_adaptor.cpp new file mode 100644 index 0000000..bf796f1 --- /dev/null +++ b/src/memory/experimental_reservation_aware_resource_adaptor.cpp @@ -0,0 +1,85 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. + * All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#include + +#include + +namespace cucascade { +namespace memory { +namespace experimental { + +reservation_aware_resource_adaptor::reservation_aware_resource_adaptor( + any_device_resource primary_mr, std::int64_t limit) + : shared_base(::cuda::mr::make_shared_resource(std::move(primary_mr), limit)) +{ +} + +std::int64_t reservation_aware_resource_adaptor::limit() const noexcept { return get().limit(); } + +void reservation_aware_resource_adaptor::set_limit(std::int64_t limit) noexcept +{ + get().set_limit(limit); +} + +std::int64_t reservation_aware_resource_adaptor::current_allocated() const noexcept +{ + return get().current_allocated(); +} + +std::int64_t reservation_aware_resource_adaptor::total_reserved() const noexcept +{ + return get().total_reserved(); +} + +std::int64_t reservation_aware_resource_adaptor::available() const noexcept +{ + return get().available(); +} + +memory_record reservation_aware_resource_adaptor::get_main_record() const +{ + return get().get_main_record(); +} + +rmm::device_async_resource_ref reservation_aware_resource_adaptor::get_upstream_resource() + const noexcept +{ + return rmm::device_async_resource_ref{ + const_cast(get().get_upstream_resource())}; +} + +memory_reservation::memory_reservation(reservation_aware_resource_adaptor const& adaptor, + std::size_t granted, + std::size_t overbooking) + : shared_base{::cuda::mr::make_shared_resource( + adaptor, detail::safe_cast(granted), overbooking)} +{ +} + +memory_reservation reservation_aware_resource_adaptor::reserve(std::size_t size, + allow_overbooking overbooking_policy) +{ + auto const [granted, overbooking] = + get().reserve(size, overbooking_policy == allow_overbooking::YES); + return memory_reservation{*this, granted, overbooking}; +} + +} // namespace experimental +} // namespace memory +} // namespace cucascade diff --git a/test/CMakeLists.txt b/test/CMakeLists.txt index 03de372..1d69b74 100644 --- a/test/CMakeLists.txt +++ b/test/CMakeLists.txt @@ -31,6 +31,7 @@ if(NOT CUCASCADE_TOPOLOGY_ONLY) # Memory tests memory/test_memory_reservation_manager.cpp memory/test_reservation_aware_resource_adaptor.cpp + memory/test_experimental_reservation_aware_resource_adaptor.cpp memory/test_small_pinned_host_memory_resource.cpp memory/test_gpu_kernels.cu # Data tests diff --git a/test/memory/test_experimental_reservation_aware_resource_adaptor.cpp b/test/memory/test_experimental_reservation_aware_resource_adaptor.cpp new file mode 100644 index 0000000..3ed5457 --- /dev/null +++ b/test/memory/test_experimental_reservation_aware_resource_adaptor.cpp @@ -0,0 +1,292 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. + * All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +/** + * Test Tags: + * [experimental_reservation_aware] - experimental reservation-aware adaptor + * [gpu] - requires a CUDA device + */ + +#include + +#include +#include +#include +#include +#include +#include + +#include +#include + +#include + +#include +#include +#include +#include +#include + +using cucascade::memory::experimental::allow_overbooking; +using cucascade::memory::experimental::memory_reservation; +using cucascade::memory::experimental::reservation_aware_resource_adaptor; + +namespace { + +bool has_cuda_device() +{ + int device_count = 0; + return cudaGetDeviceCount(&device_count) == cudaSuccess && device_count > 0; +} + +void synchronize_pool(rmm::cuda_stream_pool& pool) +{ + for (std::size_t i = 0; i < pool.get_pool_size(); ++i) { + pool.get_stream(i).synchronize(); + } +} + +constexpr std::int64_t limit = 1 << 20; + +} // namespace + +TEST_CASE("Reserve moves bytes from available to reserved", "[experimental_reservation_aware][gpu]") +{ + if (!has_cuda_device()) { return; } + + reservation_aware_resource_adaptor adaptor{ + ::cuda::mr::any_resource<::cuda::mr::device_accessible>{rmm::mr::cuda_memory_resource{}}, + limit}; + + REQUIRE(adaptor.available() == limit); + + REQUIRE_NOTHROW(std::ignore = adaptor.reserve(0, allow_overbooking::NO)); + + auto res = adaptor.reserve(1024, allow_overbooking::NO); + CHECK(adaptor == res.adaptor()); + CHECK(res.overbooking() == 0); + CHECK(res.balance() == 1024); + CHECK(adaptor.total_reserved() == 1024); + CHECK(adaptor.available() == limit - 1024); +} + +TEST_CASE("Allocating keeps available unchanged", "[experimental_reservation_aware][gpu]") +{ + if (!has_cuda_device()) { return; } + + rmm::cuda_stream_view stream{rmm::cuda_stream_default}; + reservation_aware_resource_adaptor adaptor{ + ::cuda::mr::any_resource<::cuda::mr::device_accessible>{rmm::mr::cuda_memory_resource{}}, + limit}; + + auto res = adaptor.reserve(1024, allow_overbooking::NO); + + { + rmm::device_buffer buf1{256, stream, res}; + CHECK(res.balance() == 768); + CHECK(adaptor.total_reserved() == 768); + CHECK(adaptor.current_allocated() == 256); + CHECK(adaptor.available() == limit - 1024); + + rmm::device_buffer buf2{512, stream, res}; + CHECK(res.balance() == 256); + CHECK(adaptor.total_reserved() == 256); + CHECK(adaptor.current_allocated() == 768); + CHECK(adaptor.available() == limit - 1024); + } + + CHECK(res.balance() == 1024); + CHECK(adaptor.current_allocated() == 0); + CHECK(adaptor.available() == limit - 1024); + stream.synchronize(); +} + +TEST_CASE("Exceeding the grant throws", "[experimental_reservation_aware][gpu]") +{ + if (!has_cuda_device()) { return; } + + rmm::cuda_stream_view stream{rmm::cuda_stream_default}; + reservation_aware_resource_adaptor adaptor{ + ::cuda::mr::any_resource<::cuda::mr::device_accessible>{rmm::mr::cuda_memory_resource{}}, + limit}; + + auto res = adaptor.reserve(1024, allow_overbooking::NO); + REQUIRE_THROWS_AS((rmm::device_buffer{2048, stream, res}), rmm::out_of_memory); + CHECK(res.balance() == 1024); + CHECK(adaptor.current_allocated() == 0); + stream.synchronize(); +} + +TEST_CASE("Zero-sized reservation throws on first byte", "[experimental_reservation_aware][gpu]") +{ + if (!has_cuda_device()) { return; } + + rmm::cuda_stream_view stream{rmm::cuda_stream_default}; + reservation_aware_resource_adaptor adaptor{ + ::cuda::mr::any_resource<::cuda::mr::device_accessible>{rmm::mr::cuda_memory_resource{}}, + limit}; + + auto res = adaptor.reserve(static_cast(2 * limit), allow_overbooking::NO); + CHECK(res.balance() == 0); + CHECK(res.overbooking() == static_cast(limit)); + REQUIRE_THROWS_AS((rmm::device_buffer{1, stream, res}), rmm::out_of_memory); + stream.synchronize(); +} + +TEST_CASE("Overbooking is granted when allowed", "[experimental_reservation_aware][gpu]") +{ + if (!has_cuda_device()) { return; } + + reservation_aware_resource_adaptor adaptor{ + ::cuda::mr::any_resource<::cuda::mr::device_accessible>{rmm::mr::cuda_memory_resource{}}, + limit}; + + auto res = adaptor.reserve(static_cast(2 * limit), allow_overbooking::YES); + CHECK(res.balance() == static_cast(2 * limit)); + CHECK(res.overbooking() == static_cast(limit)); + CHECK(adaptor.available() == -limit); +} + +TEST_CASE("Destruction refunds the unused balance", "[experimental_reservation_aware][gpu]") +{ + if (!has_cuda_device()) { return; } + + reservation_aware_resource_adaptor adaptor{ + ::cuda::mr::any_resource<::cuda::mr::device_accessible>{rmm::mr::cuda_memory_resource{}}, + limit}; + + { + auto res = adaptor.reserve(1024, allow_overbooking::NO); + CHECK(adaptor.total_reserved() == 1024); + } + CHECK(adaptor.total_reserved() == 0); + CHECK(adaptor.available() == limit); +} + +TEST_CASE("Buffer outlives the reserving scope", "[experimental_reservation_aware][gpu]") +{ + if (!has_cuda_device()) { return; } + + rmm::cuda_stream_view stream{rmm::cuda_stream_default}; + reservation_aware_resource_adaptor adaptor{ + ::cuda::mr::any_resource<::cuda::mr::device_accessible>{rmm::mr::cuda_memory_resource{}}, + limit}; + + { + auto buf = [&] { + auto res = adaptor.reserve(1024, allow_overbooking::NO); + return rmm::device_buffer{512, stream, res}; + }(); + + auto mr = buf.memory_resource(); + auto* reservation = ::cuda::mr::resource_cast(&mr); + REQUIRE(reservation != nullptr); + CHECK(reservation->balance() == 512); + + CHECK(adaptor.current_allocated() == 512); + CHECK(adaptor.total_reserved() == 512); + CHECK(adaptor.available() == limit - 1024); + + REQUIRE_THROWS_AS(buf.resize(2048, stream), rmm::out_of_memory); + } + + CHECK(adaptor.current_allocated() == 0); + CHECK(adaptor.total_reserved() == 0); + CHECK(adaptor.available() == limit); + stream.synchronize(); +} + +TEST_CASE("Main memory record tracks allocations", "[experimental_reservation_aware][gpu]") +{ + if (!has_cuda_device()) { return; } + + rmm::cuda_stream_view stream{rmm::cuda_stream_default}; + reservation_aware_resource_adaptor adaptor{ + ::cuda::mr::any_resource<::cuda::mr::device_accessible>{rmm::mr::cuda_memory_resource{}}, + limit}; + + auto res = adaptor.reserve(1024, allow_overbooking::NO); + { + rmm::device_buffer buf{256, stream, res}; + auto record = adaptor.get_main_record(); + CHECK(record.current == 256); + CHECK(record.total == 256); + CHECK(record.peak == 256); + CHECK(record.max == 256); + CHECK(record.num_current_allocs == 1); + CHECK(record.num_total_allocs == 1); + } + + auto record = adaptor.get_main_record(); + CHECK(record.current == 0); + CHECK(record.total == 256); + CHECK(record.peak == 256); + CHECK(record.num_current_allocs == 0); + CHECK(record.num_total_allocs == 1); + stream.synchronize(); +} + +TEST_CASE("Concurrent allocations share one reservation", "[experimental_reservation_aware][gpu]") +{ + if (!has_cuda_device()) { return; } + + constexpr std::size_t num_buffers = 100; + constexpr std::size_t max_buffer_size = 1024; + constexpr std::size_t num_threads = 2; + constexpr std::size_t grant = num_buffers * max_buffer_size; + + reservation_aware_resource_adaptor adaptor{ + ::cuda::mr::any_resource<::cuda::mr::device_accessible>{rmm::mr::cuda_memory_resource{}}, + limit}; + + std::mt19937 rng{42}; + std::uniform_int_distribution dist{0, max_buffer_size}; + std::vector sizes(num_buffers); + std::generate(sizes.begin(), sizes.end(), [&] { return dist(rng); }); + auto const total = std::accumulate(sizes.begin(), sizes.end(), std::size_t{0}); + + auto res = adaptor.reserve(grant, allow_overbooking::NO); + REQUIRE(res.balance() == grant); + + rmm::cuda_stream_pool pool{4, rmm::cuda_stream::flags::non_blocking}; + std::vector buffers(num_buffers); + std::vector> workers; + workers.reserve(num_threads); + for (std::size_t tid = 0; tid < num_threads; ++tid) { + workers.push_back(std::async(std::launch::async, [&, tid] { + for (std::size_t i = tid; i < num_buffers; i += num_threads) { + auto alloc_stream = pool.get_stream(i % pool.get_pool_size()); + buffers[i] = rmm::device_buffer{sizes[i], alloc_stream, res}; + } + })); + } + for (auto& worker : workers) { + REQUIRE_NOTHROW(worker.get()); + } + + CHECK(res.balance() == grant - total); + CHECK(adaptor.total_reserved() == static_cast(grant - total)); + CHECK(adaptor.current_allocated() == static_cast(total)); + CHECK(adaptor.available() == limit - static_cast(grant)); + + buffers.clear(); + CHECK(res.balance() == grant); + CHECK(adaptor.current_allocated() == 0); + + synchronize_pool(pool); +} From 5eba52e4fd85ea8e866bf0a3816cae85dcf3052a Mon Sep 17 00:00:00 2001 From: niranda perera Date: Wed, 5 Aug 2026 10:54:39 -0700 Subject: [PATCH 2/6] refactor Signed-off-by: niranda perera --- .../experimental/memory_reservation.hpp | 262 ++++++++++++++++++ .../reservation_aware_resource_adaptor.hpp | 236 +--------------- ...tal_reservation_aware_resource_adaptor.cpp | 2 +- ...tal_reservation_aware_resource_adaptor.cpp | 11 +- 4 files changed, 277 insertions(+), 234 deletions(-) create mode 100644 include/cucascade/memory/experimental/memory_reservation.hpp diff --git a/include/cucascade/memory/experimental/memory_reservation.hpp b/include/cucascade/memory/experimental/memory_reservation.hpp new file mode 100644 index 0000000..b1d489b --- /dev/null +++ b/include/cucascade/memory/experimental/memory_reservation.hpp @@ -0,0 +1,262 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. + * All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#pragma once + +#include + +#include +#include +#include +#include +#include + +namespace cucascade { +namespace memory { +namespace experimental { +namespace detail { + +/** + * @brief Shared state of a memory reservation. + * + * Satisfies the `cuda::mr::resource` concept so it can be a `cuda::mr::shared_resource`. + * Allocating moves bytes from the adaptor's reserved counter to its allocated counter; + * the unspent balance is refunded only when the last reference dies. + * + * @tparam Adaptor Always `reservation_aware_resource_adaptor`. + * + * The adaptor is a template parameter only so that the reservation can hold one by + * value: `reservation_aware_resource_adaptor` is defined in terms of its impl, so it is + * still incomplete at the friend declaration site, and only a dependent member type + * defers the completeness requirement to instantiation time. + */ +template + requires std::same_as +class memory_reservation_impl { + public: + memory_reservation_impl(Adaptor adaptor, std::int64_t grant, std::size_t overbooking) + : adaptor_{std::move(adaptor)}, grant_{grant}, overbooking_{overbooking}, balance_{grant} + { + } + + ~memory_reservation_impl() + { + adaptor_->total_reserved_.sub(balance(), std::memory_order_acq_rel); + } + + memory_reservation_impl(memory_reservation_impl const&) = delete; + memory_reservation_impl(memory_reservation_impl&&) = delete; + memory_reservation_impl& operator=(memory_reservation_impl const&) = delete; + memory_reservation_impl& operator=(memory_reservation_impl&&) = delete; + + [[nodiscard]] std::int64_t grant() const noexcept { return grant_; } + + [[nodiscard]] std::size_t overbooking() const noexcept { return overbooking_; } + + [[nodiscard]] std::int64_t balance() const noexcept + { + return balance_.load(std::memory_order_acquire); + } + + void* allocate(::cuda::stream_ref stream, + std::size_t bytes, + std::size_t alignment = rmm::CUDA_ALLOCATION_ALIGNMENT) + { + auto const amount = safe_cast(bytes); + draw_down_res(amount); + void* ptr = nullptr; + try { + ptr = adaptor_->allocate(stream, bytes, alignment); + } catch (...) { + balance_.fetch_add(amount, std::memory_order_acq_rel); + throw; + } + adaptor_->total_reserved_.sub(amount, std::memory_order_acq_rel); + return ptr; + } + + void deallocate(::cuda::stream_ref stream, + void* ptr, + std::size_t bytes, + std::size_t alignment = rmm::CUDA_ALLOCATION_ALIGNMENT) noexcept + { + auto const amount = safe_cast(bytes); + balance_.fetch_add(amount, std::memory_order_acq_rel); + adaptor_->total_reserved_.add(amount, std::memory_order_acq_rel); + adaptor_->deallocate(stream, ptr, bytes, alignment); + } + + void* allocate_sync(std::size_t bytes, std::size_t alignment = rmm::CUDA_ALLOCATION_ALIGNMENT) + { + auto* ptr = allocate(adaptor_->sync_stream_, bytes, alignment); + adaptor_->sync_stream_.synchronize(); + return ptr; + } + + void deallocate_sync(void* ptr, + std::size_t bytes, + std::size_t alignment = rmm::CUDA_ALLOCATION_ALIGNMENT) noexcept + { + deallocate(adaptor_->sync_stream_, ptr, bytes, alignment); + adaptor_->sync_stream_.synchronize_no_throw(); + } + + [[nodiscard]] bool operator==(memory_reservation_impl const& other) const noexcept + { + return this == std::addressof(other); + } + + friend void get_property(memory_reservation_impl const&, ::cuda::mr::device_accessible) noexcept + { + } + + [[nodiscard]] Adaptor const& adaptor() const noexcept { return adaptor_; } + + private: + void draw_down_res(std::int64_t bytes) + { + auto balance = balance_.load(std::memory_order_relaxed); + do { + if (bytes > balance) { + CUCASCADE_FAIL("allocation of " + std::to_string(bytes) + + " bytes exceeds reservation (grant: " + std::to_string(grant_) + + ", remaining: " + std::to_string(balance) + ")", + rmm::out_of_memory); + } + } while (!balance_.compare_exchange_weak( + balance, balance - bytes, std::memory_order_acq_rel, std::memory_order_relaxed)); + } + + Adaptor adaptor_; + std::int64_t const grant_; + std::size_t const overbooking_; + std::atomic balance_; +}; + +} // namespace detail + +/** + * @brief A memory reservation that is itself a memory resource. + * + * Granted by `reservation_aware_resource_adaptor::reserve()`, a reservation holds a + * budget of bytes carved out of the adaptor's limit. It is an RMM memory resource, so + * it can be handed to cudf (or anything else taking a `rmm::device_async_resource_ref`), + * and every allocation made through it is charged against that budget. An allocation + * exceeding the remaining `balance()` throws `rmm::out_of_memory`; deallocating returns + * the bytes to the balance. + * + * @par Ownership + * + * Like `reservation_aware_resource_adaptor`, this is a `cuda::mr::shared_resource`, so + * copies share the same reservation and are interchangeable. RMM stores such a copy + * inside every buffer allocated from the reservation, which is what keeps the + * reservation alive for as long as those buffers need it to service deallocations. + * + * The unspent balance is refunded to the adaptor when the last copy dies. Reserving + * more than is allocated therefore keeps the surplus out of circulation for as long as + * any derived buffer lives, so reserve what you actually use. + * + * @code{.cpp} + * auto res = adaptor.reserve(1 << 30, allow_overbooking::NO); + * auto table = cudf::groupby(..., stream, res); + * @endcode + */ +class memory_reservation : public ::cuda::mr::shared_resource< + detail::memory_reservation_impl> { + using shared_base = ::cuda::mr::shared_resource< + detail::memory_reservation_impl>; + + public: + /// @brief The shared state of the reservation. + using impl_type = detail::memory_reservation_impl; + + /** + * @brief Equality comparison. + * + * @param other The other reservation to compare. + * @return True if both refer to the same reservation. + */ + [[nodiscard]] bool operator==(memory_reservation const& other) const noexcept + { + return get() == other.get(); + } + + /** + * @brief The number of bytes originally granted. + * + * @return The granted size in bytes. + */ + [[nodiscard]] std::size_t grant() const noexcept + { + return detail::safe_cast(get().grant()); + } + + /** + * @brief The remaining unallocated size of the reservation. + * + * @return The remaining size in bytes. + */ + [[nodiscard]] std::size_t balance() const noexcept + { + return detail::safe_cast(get().balance()); + } + + /** + * @brief The number of bytes by which the grant overbooks the adaptor's limit. + * + * Nonzero only when the reservation was granted with `allow_overbooking::YES`. The + * caller must free at least this much memory before using the reservation. + * + * @return The overbooked size in bytes. + */ + [[nodiscard]] std::size_t overbooking() const noexcept { return get().overbooking(); } + + /** + * @brief The adaptor that granted the reservation. + * + * @return The adaptor. + */ + [[nodiscard]] reservation_aware_resource_adaptor const& adaptor() const noexcept + { + return get().adaptor(); + } + + private: + friend class reservation_aware_resource_adaptor; + + /** + * @brief Construct from an already-granted reservation. + * + * Private so that only `reservation_aware_resource_adaptor` can grant reservations. + * The reservation holds a copy of the adaptor, so the adaptor stays alive for as + * long as any buffer allocated from the reservation needs it. + * + * @param adaptor The adaptor that granted the reservation. + * @param granted The number of bytes granted. + * @param overbooking The number of bytes by which @p granted overbooks the limit. + */ + memory_reservation(reservation_aware_resource_adaptor const& adaptor, + std::size_t granted, + std::size_t overbooking); +}; + +static_assert(::cuda::mr::resource_with); + +} // namespace experimental +} // namespace memory +} // namespace cucascade diff --git a/include/cucascade/memory/experimental/reservation_aware_resource_adaptor.hpp b/include/cucascade/memory/experimental/reservation_aware_resource_adaptor.hpp index 3957d31..a927615 100644 --- a/include/cucascade/memory/experimental/reservation_aware_resource_adaptor.hpp +++ b/include/cucascade/memory/experimental/reservation_aware_resource_adaptor.hpp @@ -82,10 +82,8 @@ template } } -// The adaptor is a template parameter only so that the reservation can hold one by -// value: `reservation_aware_resource_adaptor` is defined in terms of the adaptor impl -// below, so it is still incomplete here, and only a dependent member type defers the -// completeness requirement to instantiation time. +// Forward-declared so `reservation_aware_resource_adaptor_impl` can friend it. +// Defined in memory_reservation.hpp. template requires std::same_as class memory_reservation_impl; @@ -267,118 +265,6 @@ class reservation_aware_resource_adaptor_impl { rmm::cuda_stream sync_stream_{rmm::cuda_stream::flags::non_blocking}; }; -/** - * @brief Shared state of a memory reservation. - * - * Satisfies the `cuda::mr::resource` concept so it can be a `cuda::mr::shared_resource`. - * Allocating moves bytes from the adaptor's reserved counter to its allocated counter; - * the unspent balance is refunded only when the last reference dies. - * - * @tparam Adaptor Always `reservation_aware_resource_adaptor`. - */ -template - requires std::same_as -class memory_reservation_impl { - public: - memory_reservation_impl(Adaptor adaptor, std::int64_t grant, std::size_t overbooking) - : adaptor_{std::move(adaptor)}, grant_{grant}, overbooking_{overbooking}, balance_{grant} - { - } - - ~memory_reservation_impl() - { - adaptor_->total_reserved_.sub(balance(), std::memory_order_acq_rel); - } - - memory_reservation_impl(memory_reservation_impl const&) = delete; - memory_reservation_impl(memory_reservation_impl&&) = delete; - memory_reservation_impl& operator=(memory_reservation_impl const&) = delete; - memory_reservation_impl& operator=(memory_reservation_impl&&) = delete; - - [[nodiscard]] std::int64_t grant() const noexcept { return grant_; } - - [[nodiscard]] std::size_t overbooking() const noexcept { return overbooking_; } - - [[nodiscard]] std::int64_t balance() const noexcept - { - return balance_.load(std::memory_order_acquire); - } - - void* allocate(::cuda::stream_ref stream, - std::size_t bytes, - std::size_t alignment = rmm::CUDA_ALLOCATION_ALIGNMENT) - { - auto const amount = safe_cast(bytes); - draw_down_res(amount); - void* ptr = nullptr; - try { - ptr = adaptor_->allocate(stream, bytes, alignment); - } catch (...) { - balance_.fetch_add(amount, std::memory_order_acq_rel); - throw; - } - adaptor_->total_reserved_.sub(amount, std::memory_order_acq_rel); - return ptr; - } - - void deallocate(::cuda::stream_ref stream, - void* ptr, - std::size_t bytes, - std::size_t alignment = rmm::CUDA_ALLOCATION_ALIGNMENT) noexcept - { - auto const amount = safe_cast(bytes); - balance_.fetch_add(amount, std::memory_order_acq_rel); - adaptor_->total_reserved_.add(amount, std::memory_order_acq_rel); - adaptor_->deallocate(stream, ptr, bytes, alignment); - } - - void* allocate_sync(std::size_t bytes, std::size_t alignment = rmm::CUDA_ALLOCATION_ALIGNMENT) - { - auto* ptr = allocate(adaptor_->sync_stream_, bytes, alignment); - adaptor_->sync_stream_.synchronize(); - return ptr; - } - - void deallocate_sync(void* ptr, - std::size_t bytes, - std::size_t alignment = rmm::CUDA_ALLOCATION_ALIGNMENT) noexcept - { - deallocate(adaptor_->sync_stream_, ptr, bytes, alignment); - adaptor_->sync_stream_.synchronize_no_throw(); - } - - [[nodiscard]] bool operator==(memory_reservation_impl const& other) const noexcept - { - return this == std::addressof(other); - } - - friend void get_property(memory_reservation_impl const&, ::cuda::mr::device_accessible) noexcept - { - } - - [[nodiscard]] Adaptor const& adaptor() const noexcept { return adaptor_; } - - private: - void draw_down_res(std::int64_t bytes) - { - auto balance = balance_.load(std::memory_order_relaxed); - do { - if (bytes > balance) { - CUCASCADE_FAIL("allocation of " + std::to_string(bytes) + - " bytes exceeds reservation (grant: " + std::to_string(grant_) + - ", remaining: " + std::to_string(balance) + ")", - rmm::out_of_memory); - } - } while (!balance_.compare_exchange_weak( - balance, balance - bytes, std::memory_order_acq_rel, std::memory_order_relaxed)); - } - - Adaptor adaptor_; - std::int64_t const grant_; - std::size_t const overbooking_; - std::atomic balance_; -}; - } // namespace detail /** @@ -417,12 +303,6 @@ class reservation_aware_resource_adaptor /// @brief The reference-counted handle on the shared implementation. using shared_base = ::cuda::mr::shared_resource; - /// @brief Tag this resource as device-accessible for the CCCL concept. - friend void get_property(reservation_aware_resource_adaptor const&, - ::cuda::mr::device_accessible) noexcept - { - } - /** * @brief Construct with the specified primary memory resource and limit. * @@ -523,116 +403,8 @@ class reservation_aware_resource_adaptor static_assert( ::cuda::mr::resource_with); -/** - * @brief A memory reservation that is itself a memory resource. - * - * Granted by `reservation_aware_resource_adaptor::reserve()`, a reservation holds a - * budget of bytes carved out of the adaptor's limit. It is an RMM memory resource, so - * it can be handed to cudf (or anything else taking a `rmm::device_async_resource_ref`), - * and every allocation made through it is charged against that budget. An allocation - * exceeding the remaining `balance()` throws `rmm::out_of_memory`; deallocating returns - * the bytes to the balance. - * - * @par Ownership - * - * Like `reservation_aware_resource_adaptor`, this is a `cuda::mr::shared_resource`, so - * copies share the same reservation and are interchangeable. RMM stores such a copy - * inside every buffer allocated from the reservation, which is what keeps the - * reservation alive for as long as those buffers need it to service deallocations. - * - * The unspent balance is refunded to the adaptor when the last copy dies. Reserving - * more than is allocated therefore keeps the surplus out of circulation for as long as - * any derived buffer lives, so reserve what you actually use. - * - * @code{.cpp} - * auto res = adaptor.reserve(1 << 30, allow_overbooking::NO); - * auto table = cudf::groupby(..., stream, res); - * @endcode - */ -class memory_reservation : public ::cuda::mr::shared_resource< - detail::memory_reservation_impl> { - using shared_base = ::cuda::mr::shared_resource< - detail::memory_reservation_impl>; - - public: - /// @brief The shared state of the reservation. - using impl_type = detail::memory_reservation_impl; - - /// @brief Tag this resource as device-accessible for the CCCL concept. - friend void get_property(memory_reservation const&, ::cuda::mr::device_accessible) noexcept {} - - /** - * @brief Equality comparison. - * - * @param other The other reservation to compare. - * @return True if both refer to the same reservation. - */ - [[nodiscard]] bool operator==(memory_reservation const& other) const noexcept - { - return get() == other.get(); - } - - /** - * @brief The number of bytes originally granted. - * - * @return The granted size in bytes. - */ - [[nodiscard]] std::size_t grant() const noexcept - { - return detail::safe_cast(get().grant()); - } - - /** - * @brief The remaining unallocated size of the reservation. - * - * @return The remaining size in bytes. - */ - [[nodiscard]] std::size_t balance() const noexcept - { - return detail::safe_cast(get().balance()); - } - - /** - * @brief The number of bytes by which the grant overbooks the adaptor's limit. - * - * Nonzero only when the reservation was granted with `allow_overbooking::YES`. The - * caller must free at least this much memory before using the reservation. - * - * @return The overbooked size in bytes. - */ - [[nodiscard]] std::size_t overbooking() const noexcept { return get().overbooking(); } - - /** - * @brief The adaptor that granted the reservation. - * - * @return The adaptor. - */ - [[nodiscard]] reservation_aware_resource_adaptor const& adaptor() const noexcept - { - return get().adaptor(); - } - - private: - friend class reservation_aware_resource_adaptor; - - /** - * @brief Construct from an already-granted reservation. - * - * Private so that only `reservation_aware_resource_adaptor` can grant reservations. - * The reservation holds a copy of the adaptor, so the adaptor stays alive for as - * long as any buffer allocated from the reservation needs it. - * - * @param adaptor The adaptor that granted the reservation. - * @param granted The number of bytes granted. - * @param overbooking The number of bytes by which @p granted overbooks the limit. - */ - memory_reservation(reservation_aware_resource_adaptor const& adaptor, - std::size_t granted, - std::size_t overbooking); -}; - -static_assert(::cuda::mr::resource_with); - } // namespace experimental } // namespace memory } // namespace cucascade + +#include diff --git a/src/memory/experimental_reservation_aware_resource_adaptor.cpp b/src/memory/experimental_reservation_aware_resource_adaptor.cpp index bf796f1..cb775a4 100644 --- a/src/memory/experimental_reservation_aware_resource_adaptor.cpp +++ b/src/memory/experimental_reservation_aware_resource_adaptor.cpp @@ -16,7 +16,7 @@ * limitations under the License. */ -#include +#include #include diff --git a/test/memory/test_experimental_reservation_aware_resource_adaptor.cpp b/test/memory/test_experimental_reservation_aware_resource_adaptor.cpp index 3ed5457..e9f34a4 100644 --- a/test/memory/test_experimental_reservation_aware_resource_adaptor.cpp +++ b/test/memory/test_experimental_reservation_aware_resource_adaptor.cpp @@ -73,12 +73,21 @@ TEST_CASE("Reserve moves bytes from available to reserved", "[experimental_reser ::cuda::mr::any_resource<::cuda::mr::device_accessible>{rmm::mr::cuda_memory_resource{}}, limit}; + auto check_any_resouce_conversion = [&](auto& mr_like) { + cuda::mr::any_resource any_device = mr_like; + cuda::mr::any_resource<> any_adaptor = mr_like; + CHECK(any_device == mr_like); + CHECK(any_adaptor == mr_like); + }; + check_any_resouce_conversion(adaptor); + REQUIRE(adaptor.available() == limit); REQUIRE_NOTHROW(std::ignore = adaptor.reserve(0, allow_overbooking::NO)); auto res = adaptor.reserve(1024, allow_overbooking::NO); - CHECK(adaptor == res.adaptor()); + check_any_resouce_conversion(res); + CHECK(res.overbooking() == 0); CHECK(res.balance() == 1024); CHECK(adaptor.total_reserved() == 1024); From 82ae25ce49d0d912ced4b8e2b6b15edd115c491c Mon Sep 17 00:00:00 2001 From: niranda perera Date: Wed, 5 Aug 2026 15:29:10 -0700 Subject: [PATCH 3/6] extend reservations to handle accessibility Signed-off-by: niranda perera --- .../experimental/memory_reservation.hpp | 176 +++++++++++------- .../reservation_aware_resource_adaptor.hpp | 95 +++++++--- src/memory/CMakeLists.txt | 3 +- .../experimental/memory_reservation.cpp | 149 +++++++++++++++ .../reservation_aware_resource_adaptor.cpp | 102 ++++++++++ ...tal_reservation_aware_resource_adaptor.cpp | 85 --------- ...tal_reservation_aware_resource_adaptor.cpp | 145 ++++++++++----- 7 files changed, 530 insertions(+), 225 deletions(-) create mode 100644 src/memory/experimental/memory_reservation.cpp create mode 100644 src/memory/experimental/reservation_aware_resource_adaptor.cpp delete mode 100644 src/memory/experimental_reservation_aware_resource_adaptor.cpp diff --git a/include/cucascade/memory/experimental/memory_reservation.hpp b/include/cucascade/memory/experimental/memory_reservation.hpp index b1d489b..fb74a00 100644 --- a/include/cucascade/memory/experimental/memory_reservation.hpp +++ b/include/cucascade/memory/experimental/memory_reservation.hpp @@ -25,6 +25,7 @@ #include #include #include +#include namespace cucascade { namespace memory { @@ -38,16 +39,16 @@ namespace detail { * Allocating moves bytes from the adaptor's reserved counter to its allocated counter; * the unspent balance is refunded only when the last reference dies. * - * @tparam Adaptor Always `reservation_aware_resource_adaptor`. - * - * The adaptor is a template parameter only so that the reservation can hold one by - * value: `reservation_aware_resource_adaptor` is defined in terms of its impl, so it is - * still incomplete at the friend declaration site, and only a dependent member type - * defers the completeness requirement to instantiation time. + * @tparam Adaptor One of `device_adaptor`, `host_adaptor`, or `host_device_adaptor`. Its + * properties are forwarded to the reservation via `cuda::forward_property`, so a + * reservation advertises whatever the granting adaptor advertises (e.g. + * `cuda::mr::device_accessible`). That forwarding is what lets `memory_reservation` + * project this impl into an erased resource of the matching accessibility. */ template - requires std::same_as -class memory_reservation_impl { + requires reservation_adaptor +class memory_reservation_impl + : public ::cuda::forward_property, Adaptor> { public: memory_reservation_impl(Adaptor adaptor, std::int64_t grant, std::size_t overbooking) : adaptor_{std::move(adaptor)}, grant_{grant}, overbooking_{overbooking}, balance_{grant} @@ -121,9 +122,7 @@ class memory_reservation_impl { return this == std::addressof(other); } - friend void get_property(memory_reservation_impl const&, ::cuda::mr::device_accessible) noexcept - { - } + [[nodiscard]] Adaptor const& upstream_resource() const noexcept { return adaptor_; } [[nodiscard]] Adaptor const& adaptor() const noexcept { return adaptor_; } @@ -148,42 +147,102 @@ class memory_reservation_impl { std::atomic balance_; }; +/** + * @brief The reference-counted handle on a reservation's shared state. + */ +template +using reservation_handle = ::cuda::mr::shared_resource>; + } // namespace detail /** - * @brief A memory reservation that is itself a memory resource. + * @brief A memory reservation, independent of the accessibility of its memory. * * Granted by `reservation_aware_resource_adaptor::reserve()`, a reservation holds a - * budget of bytes carved out of the adaptor's limit. It is an RMM memory resource, so - * it can be handed to cudf (or anything else taking a `rmm::device_async_resource_ref`), - * and every allocation made through it is charged against that budget. An allocation - * exceeding the remaining `balance()` throws `rmm::out_of_memory`; deallocating returns - * the bytes to the balance. + * budget of bytes carved out of the adaptor's limit. Every allocation made through it is + * charged against that budget: an allocation exceeding the remaining `balance()` throws + * `rmm::out_of_memory`, and deallocating returns the bytes to the balance. * - * @par Ownership + * This is a single type regardless of where the memory lives. Internally it holds one of + * three shared states, chosen by the granting adaptor's upstream, and its own type says + * nothing about accessibility. Query `accessibility()` to find out. * - * Like `reservation_aware_resource_adaptor`, this is a `cuda::mr::shared_resource`, so - * copies share the same reservation and are interchangeable. RMM stores such a copy - * inside every buffer allocated from the reservation, which is what keeps the - * reservation alive for as long as those buffers need it to service deallocations. + * @par Handing a reservation to cudf or RMM * - * The unspent balance is refunded to the adaptor when the last copy dies. Reserving - * more than is allocated therefore keeps the surplus out of circulation for as long as - * any derived buffer lives, so reserve what you actually use. + * A reservation is not itself a memory resource. Project it into an erased resource of + * the accessibility you need, and pass that: * * @code{.cpp} - * auto res = adaptor.reserve(1 << 30, allow_overbooking::NO); - * auto table = cudf::groupby(..., stream, res); + * auto res = adaptor.reserve(1 << 30, allow_overbooking::NO); + * auto table = cudf::groupby(..., stream, res.as_device()); * @endcode + * + * `as_device()`, `as_host()`, and `as_host_device()` each throw + * `cucascade::logic_error` when the reservation's memory does not have the requested + * accessibility. Use `accessibility()`, `is_device_accessible()`, or + * `is_host_accessible()` to branch without exceptions. + * + * @par Ownership + * + * Copies of a reservation share the same shared state and are interchangeable. The + * handles returned by the `as_*()` methods own a reference to that same state, so a + * buffer allocated through one keeps the reservation alive for as long as it needs to + * service its deallocation, even after every `memory_reservation` copy is gone. + * + * The unspent balance is refunded to the adaptor when the last reference dies. Reserving + * more than is allocated therefore keeps the surplus out of circulation for as long as + * any derived buffer lives, so reserve what you actually use. */ -class memory_reservation : public ::cuda::mr::shared_resource< - detail::memory_reservation_impl> { - using shared_base = ::cuda::mr::shared_resource< - detail::memory_reservation_impl>; - +class memory_reservation { public: - /// @brief The shared state of the reservation. - using impl_type = detail::memory_reservation_impl; + /** + * @brief The accessibility of the memory this reservation draws on. + * + * @return The accessibility, fixed at grant time by the adaptor's upstream. + */ + [[nodiscard]] reservation_accessibility accessibility() const noexcept; + + /** + * @brief Whether `as_device()` will succeed. + * + * @return True when the reservation's memory is device accessible. + */ + [[nodiscard]] bool is_device_accessible() const noexcept; + + /** + * @brief Whether `as_host()` will succeed. + * + * @return True when the reservation's memory is host accessible. + */ + [[nodiscard]] bool is_host_accessible() const noexcept; + + /** + * @brief Project the reservation into a device-accessible memory resource. + * + * The returned handle owns a reference to the reservation's shared state, so it may + * outlive this object. + * + * @return An erased resource advertising `cuda::mr::device_accessible`. + * @throws cucascade::logic_error if the reservation's memory is not device accessible. + */ + [[nodiscard]] any_device_resource as_device() const; + + /** + * @brief Project the reservation into a host-accessible memory resource. + * + * @return An erased resource advertising `cuda::mr::host_accessible`. + * @throws cucascade::logic_error if the reservation's memory is not host accessible. + */ + [[nodiscard]] any_host_resource as_host() const; + + /** + * @brief Project the reservation into a host- and device-accessible memory resource. + * + * @return An erased resource advertising both accessibility properties. + * @throws cucascade::logic_error unless the reservation's memory is both host and + * device accessible. + */ + [[nodiscard]] any_host_device_resource as_host_device() const; /** * @brief Equality comparison. @@ -191,30 +250,21 @@ class memory_reservation : public ::cuda::mr::shared_resource< * @param other The other reservation to compare. * @return True if both refer to the same reservation. */ - [[nodiscard]] bool operator==(memory_reservation const& other) const noexcept - { - return get() == other.get(); - } + [[nodiscard]] bool operator==(memory_reservation const& other) const noexcept; /** * @brief The number of bytes originally granted. * * @return The granted size in bytes. */ - [[nodiscard]] std::size_t grant() const noexcept - { - return detail::safe_cast(get().grant()); - } + [[nodiscard]] std::size_t grant() const noexcept; /** * @brief The remaining unallocated size of the reservation. * * @return The remaining size in bytes. */ - [[nodiscard]] std::size_t balance() const noexcept - { - return detail::safe_cast(get().balance()); - } + [[nodiscard]] std::size_t balance() const noexcept; /** * @brief The number of bytes by which the grant overbooks the adaptor's limit. @@ -224,38 +274,28 @@ class memory_reservation : public ::cuda::mr::shared_resource< * * @return The overbooked size in bytes. */ - [[nodiscard]] std::size_t overbooking() const noexcept { return get().overbooking(); } - - /** - * @brief The adaptor that granted the reservation. - * - * @return The adaptor. - */ - [[nodiscard]] reservation_aware_resource_adaptor const& adaptor() const noexcept - { - return get().adaptor(); - } + [[nodiscard]] std::size_t overbooking() const noexcept; private: + template friend class reservation_aware_resource_adaptor; + /// @brief The shared state, one alternative per accessibility. + using handle_variant = std::variant, + detail::reservation_handle, + detail::reservation_handle>; + /** - * @brief Construct from an already-granted reservation. + * @brief Construct from an already-granted shared state. * * Private so that only `reservation_aware_resource_adaptor` can grant reservations. - * The reservation holds a copy of the adaptor, so the adaptor stays alive for as - * long as any buffer allocated from the reservation needs it. * - * @param adaptor The adaptor that granted the reservation. - * @param granted The number of bytes granted. - * @param overbooking The number of bytes by which @p granted overbooks the limit. + * @param handle The shared state of the granted reservation. */ - memory_reservation(reservation_aware_resource_adaptor const& adaptor, - std::size_t granted, - std::size_t overbooking); -}; + explicit memory_reservation(handle_variant handle) : handle_{std::move(handle)} {} -static_assert(::cuda::mr::resource_with); + handle_variant handle_; +}; } // namespace experimental } // namespace memory diff --git a/include/cucascade/memory/experimental/reservation_aware_resource_adaptor.hpp b/include/cucascade/memory/experimental/reservation_aware_resource_adaptor.hpp index a927615..9307f03 100644 --- a/include/cucascade/memory/experimental/reservation_aware_resource_adaptor.hpp +++ b/include/cucascade/memory/experimental/reservation_aware_resource_adaptor.hpp @@ -42,10 +42,31 @@ namespace cucascade { namespace memory { namespace experimental { +template class reservation_aware_resource_adaptor; class memory_reservation; using any_device_resource = ::cuda::mr::any_resource<::cuda::mr::device_accessible>; +using any_host_resource = ::cuda::mr::any_resource<::cuda::mr::host_accessible>; +using any_host_device_resource = + ::cuda::mr::any_resource<::cuda::mr::device_accessible, ::cuda::mr::host_accessible>; + +using device_adaptor = reservation_aware_resource_adaptor; +using host_adaptor = reservation_aware_resource_adaptor; +using host_device_adaptor = reservation_aware_resource_adaptor; + +/** + * @brief The accessibility of the memory a reservation draws on. + * + * Determined by the upstream of the granting adaptor and fixed for the lifetime of the + * reservation. Selects which of `memory_reservation::as_device()`, `as_host()`, and + * `as_host_device()` are valid. + */ +enum class reservation_accessibility : std::uint8_t { + DEVICE, ///< Device accessible only. + HOST, ///< Host accessible only. + HOST_DEVICE, ///< Both host and device accessible. +}; /** * @brief Snapshot of the adaptor's main (non-scoped) allocation accounting. @@ -82,19 +103,34 @@ template } } +/** + * @brief The adaptor instantiations that may grant a reservation. + */ +template +concept reservation_adaptor = std::same_as || std::same_as || + std::same_as; + // Forward-declared so `reservation_aware_resource_adaptor_impl` can friend it. // Defined in memory_reservation.hpp. template - requires std::same_as + requires reservation_adaptor class memory_reservation_impl; /** * @brief Shared state of a reservation-aware resource adaptor. * - * Owns an upstream device resource and tracks a single main memory record (no scoped + * Owns an upstream resource and tracks a single main memory record (no scoped * records). Reservations are granted against a runtime-adjustable limit. + * + * @tparam Upstream The upstream memory resource type. Its properties are forwarded to + * the impl via `cuda::forward_property`, so any tag advertised by `Upstream` + * (e.g. `cuda::mr::device_accessible`) is visible on the impl and, transitively, on + * the wrapping `shared_resource`. */ -class reservation_aware_resource_adaptor_impl { +template + requires ::cuda::mr::resource +class reservation_aware_resource_adaptor_impl + : public ::cuda::forward_property, Upstream> { public: /** * @brief Construct with a primary memory resource and a memory limit. @@ -102,7 +138,7 @@ class reservation_aware_resource_adaptor_impl { * @param upstream_mr The primary memory resource (moved in). * @param limit Maximum number of bytes that may be allocated and reserved. */ - reservation_aware_resource_adaptor_impl(any_device_resource upstream_mr, std::int64_t limit) + reservation_aware_resource_adaptor_impl(Upstream upstream_mr, std::int64_t limit) : upstream_mr_{std::move(upstream_mr)}, limit_{limit} { } @@ -121,10 +157,9 @@ class reservation_aware_resource_adaptor_impl { return this == std::addressof(other); } - [[nodiscard]] any_device_resource const& get_upstream_resource() const noexcept - { - return upstream_mr_; - } + [[nodiscard]] Upstream const& upstream_resource() const noexcept { return upstream_mr_; } + + [[nodiscard]] Upstream const& get_upstream_resource() const noexcept { return upstream_mr_; } [[nodiscard]] std::int64_t limit() const noexcept { @@ -225,13 +260,10 @@ class reservation_aware_resource_adaptor_impl { sync_stream_.synchronize_no_throw(); } - friend void get_property(reservation_aware_resource_adaptor_impl const&, - ::cuda::mr::device_accessible) noexcept - { - } - private: - friend class memory_reservation_impl; + friend class memory_reservation_impl; + friend class memory_reservation_impl; + friend class memory_reservation_impl; void record_allocation(std::int64_t nbytes) { @@ -249,7 +281,7 @@ class reservation_aware_resource_adaptor_impl { num_current_allocs_.sub(1, std::memory_order_acq_rel); } - any_device_resource upstream_mr_; + Upstream upstream_mr_; utils::atomic_bounded_counter num_current_allocs_{0}; utils::atomic_bounded_counter num_total_allocs_{0}; @@ -270,13 +302,21 @@ class reservation_aware_resource_adaptor_impl { /** * @brief A memory resource adaptor that only allocates through reservations. * - * This adaptor wraps a primary device memory resource and adds a memory limit with - * allocation tracking. Memory is obtained by calling `reserve()` and allocating through - * the returned `memory_reservation`. + * This adaptor wraps a primary memory resource and adds a memory limit with allocation + * tracking. Memory is obtained by calling `reserve()` and allocating through the returned + * `memory_reservation`. * * This class is copyable and shares ownership of its internal state via * `cuda::mr::shared_resource`. * + * Only three instantiations exist, one per accessibility: `device_adaptor`, + * `host_adaptor`, and `host_device_adaptor`. The upstream's accessibility is forwarded to + * the adaptor and on to every reservation it grants, so a `device_adaptor` yields + * reservations usable as device resources and nothing else. + * + * @tparam Upstream The erased upstream resource type: `any_device_resource`, + * `any_host_resource`, or `any_host_device_resource`. + * * @par Allocating without a reservation * * The adaptor is itself a memory resource, so it can be handed to cudf or an RMM @@ -294,22 +334,26 @@ class reservation_aware_resource_adaptor_impl { * Allocating through a reservation moves bytes from the second bucket to the first, * leaving `available()` unchanged; that is what makes a reservation a promise. */ +template class reservation_aware_resource_adaptor - : public ::cuda::mr::shared_resource { + : public ::cuda::mr::shared_resource> { public: /// @brief The adaptor's shared implementation. - using impl_type = detail::reservation_aware_resource_adaptor_impl; + using impl_type = detail::reservation_aware_resource_adaptor_impl; /// @brief The reference-counted handle on the shared implementation. using shared_base = ::cuda::mr::shared_resource; + /// @brief The erased upstream resource type. + using upstream_type = Upstream; + /** * @brief Construct with the specified primary memory resource and limit. * * @param upstream_mr The primary memory resource. * @param limit Maximum number of bytes that may be allocated and reserved. */ - reservation_aware_resource_adaptor(any_device_resource upstream_mr, std::int64_t limit); + reservation_aware_resource_adaptor(Upstream upstream_mr, std::int64_t limit); /** * @brief Equality comparison. @@ -319,7 +363,7 @@ class reservation_aware_resource_adaptor */ [[nodiscard]] bool operator==(reservation_aware_resource_adaptor const& other) const noexcept { - return get() == other.get(); + return this->get() == other.get(); } /** @@ -395,14 +439,11 @@ class reservation_aware_resource_adaptor /** * @brief Get a reference to the primary upstream resource. * - * @return Reference to the RMM memory resource. + * @return Reference to the erased upstream resource. */ - [[nodiscard]] rmm::device_async_resource_ref get_upstream_resource() const noexcept; + [[nodiscard]] Upstream const& get_upstream_resource() const noexcept; }; -static_assert( - ::cuda::mr::resource_with); - } // namespace experimental } // namespace memory } // namespace cucascade diff --git a/src/memory/CMakeLists.txt b/src/memory/CMakeLists.txt index 007c2f8..6686eb6 100644 --- a/src/memory/CMakeLists.txt +++ b/src/memory/CMakeLists.txt @@ -30,7 +30,8 @@ if(TARGET cucascade_objects) ${CMAKE_CURRENT_SOURCE_DIR}/numa_region_pinned_host_allocator.cpp ${CMAKE_CURRENT_SOURCE_DIR}/oom_handling_policy.cpp ${CMAKE_CURRENT_SOURCE_DIR}/reservation_aware_resource_adaptor.cpp - ${CMAKE_CURRENT_SOURCE_DIR}/experimental_reservation_aware_resource_adaptor.cpp + ${CMAKE_CURRENT_SOURCE_DIR}/experimental/memory_reservation.cpp + ${CMAKE_CURRENT_SOURCE_DIR}/experimental/reservation_aware_resource_adaptor.cpp ${CMAKE_CURRENT_SOURCE_DIR}/reservation_manager_configurator.cpp ${CMAKE_CURRENT_SOURCE_DIR}/small_pinned_host_memory_resource.cpp ${CMAKE_CURRENT_SOURCE_DIR}/stream_pool.cpp diff --git a/src/memory/experimental/memory_reservation.cpp b/src/memory/experimental/memory_reservation.cpp new file mode 100644 index 0000000..24f0f86 --- /dev/null +++ b/src/memory/experimental/memory_reservation.cpp @@ -0,0 +1,149 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. + * All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#include + +#include +#include + +namespace cucascade { +namespace memory { +namespace experimental { + +namespace { + +template +constexpr bool handle_is_device_accessible = + ::cuda::has_property; + +template +constexpr bool handle_is_host_accessible = + ::cuda::has_property; + +template +[[nodiscard]] constexpr reservation_accessibility accessibility_of() noexcept +{ + static_assert(handle_is_device_accessible || handle_is_host_accessible, + "a reservation must be accessible from somewhere"); + if constexpr (handle_is_device_accessible && handle_is_host_accessible) { + return reservation_accessibility::HOST_DEVICE; + } else if constexpr (handle_is_device_accessible) { + return reservation_accessibility::DEVICE; + } else { + return reservation_accessibility::HOST; + } +} + +} // namespace + +// Each variant alternative must forward the accessibility of the adaptor it was granted from, +// so that the corresponding as_*() projection is well-formed. +static_assert(::cuda::mr::resource_with, + ::cuda::mr::device_accessible>); +static_assert( + ::cuda::mr::resource_with, ::cuda::mr::host_accessible>); +static_assert(::cuda::mr::resource_with, + ::cuda::mr::device_accessible, + ::cuda::mr::host_accessible>); + +reservation_accessibility memory_reservation::accessibility() const noexcept +{ + return std::visit( + [](auto const& handle) { return accessibility_of>(); }, + handle_); +} + +bool memory_reservation::is_device_accessible() const noexcept +{ + auto const access = accessibility(); + return access == reservation_accessibility::DEVICE || + access == reservation_accessibility::HOST_DEVICE; +} + +bool memory_reservation::is_host_accessible() const noexcept +{ + auto const access = accessibility(); + return access == reservation_accessibility::HOST || + access == reservation_accessibility::HOST_DEVICE; +} + +any_device_resource memory_reservation::as_device() const +{ + return std::visit( + [](auto const& handle) -> any_device_resource { + if constexpr (handle_is_device_accessible>) { + return any_device_resource{handle}; + } else { + CUCASCADE_FAIL("reservation memory is not device accessible"); + } + }, + handle_); +} + +any_host_resource memory_reservation::as_host() const +{ + return std::visit( + [](auto const& handle) -> any_host_resource { + if constexpr (handle_is_host_accessible>) { + return any_host_resource{handle}; + } else { + CUCASCADE_FAIL("reservation memory is not host accessible"); + } + }, + handle_); +} + +any_host_device_resource memory_reservation::as_host_device() const +{ + return std::visit( + [](auto const& handle) -> any_host_device_resource { + using handle_t = std::remove_cvref_t; + if constexpr (handle_is_device_accessible && handle_is_host_accessible) { + return any_host_device_resource{handle}; + } else { + CUCASCADE_FAIL("reservation memory is not both host and device accessible"); + } + }, + handle_); +} + +bool memory_reservation::operator==(memory_reservation const& other) const noexcept +{ + return handle_ == other.handle_; +} + +std::size_t memory_reservation::grant() const noexcept +{ + return std::visit( + [](auto const& handle) { return detail::safe_cast(handle->grant()); }, handle_); +} + +std::size_t memory_reservation::balance() const noexcept +{ + return std::visit( + [](auto const& handle) { return detail::safe_cast(handle->balance()); }, handle_); +} + +std::size_t memory_reservation::overbooking() const noexcept +{ + return std::visit([](auto const& handle) { return handle->overbooking(); }, handle_); +} + +} // namespace experimental +} // namespace memory +} // namespace cucascade diff --git a/src/memory/experimental/reservation_aware_resource_adaptor.cpp b/src/memory/experimental/reservation_aware_resource_adaptor.cpp new file mode 100644 index 0000000..58982ab --- /dev/null +++ b/src/memory/experimental/reservation_aware_resource_adaptor.cpp @@ -0,0 +1,102 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. + * All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#include + +#include + +namespace cucascade { +namespace memory { +namespace experimental { + +template +reservation_aware_resource_adaptor::reservation_aware_resource_adaptor( + Upstream primary_mr, std::int64_t limit) + : shared_base(::cuda::mr::make_shared_resource(std::move(primary_mr), limit)) +{ +} + +template +std::int64_t reservation_aware_resource_adaptor::limit() const noexcept +{ + return this->get().limit(); +} + +template +void reservation_aware_resource_adaptor::set_limit(std::int64_t limit) noexcept +{ + this->get().set_limit(limit); +} + +template +std::int64_t reservation_aware_resource_adaptor::current_allocated() const noexcept +{ + return this->get().current_allocated(); +} + +template +std::int64_t reservation_aware_resource_adaptor::total_reserved() const noexcept +{ + return this->get().total_reserved(); +} + +template +std::int64_t reservation_aware_resource_adaptor::available() const noexcept +{ + return this->get().available(); +} + +template +memory_record reservation_aware_resource_adaptor::get_main_record() const +{ + return this->get().get_main_record(); +} + +template +Upstream const& reservation_aware_resource_adaptor::get_upstream_resource() const noexcept +{ + return this->get().get_upstream_resource(); +} + +template +memory_reservation reservation_aware_resource_adaptor::reserve( + std::size_t size, allow_overbooking overbooking_policy) +{ + auto const [granted, overbooking] = + this->get().reserve(size, overbooking_policy == allow_overbooking::YES); + + using impl_t = detail::memory_reservation_impl>; + return memory_reservation{ + memory_reservation::handle_variant{::cuda::mr::make_shared_resource( + *this, detail::safe_cast(granted), overbooking)}}; +} + +template class reservation_aware_resource_adaptor; +template class reservation_aware_resource_adaptor; +template class reservation_aware_resource_adaptor; + +// Each adaptor must forward the accessibility of the upstream it was instantiated with. +static_assert(::cuda::mr::resource_with); +static_assert(::cuda::mr::resource_with); +static_assert(::cuda::mr::resource_with); + +} // namespace experimental +} // namespace memory +} // namespace cucascade diff --git a/src/memory/experimental_reservation_aware_resource_adaptor.cpp b/src/memory/experimental_reservation_aware_resource_adaptor.cpp deleted file mode 100644 index cb775a4..0000000 --- a/src/memory/experimental_reservation_aware_resource_adaptor.cpp +++ /dev/null @@ -1,85 +0,0 @@ -/* - * SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. - * All rights reserved. - * SPDX-License-Identifier: Apache-2.0 - * - * Licensed under the Apache License, Version 2.0 (the "License"); - * you may not use this file except in compliance with the License. - * You may obtain a copy of the License at - * - * http://www.apache.org/licenses/LICENSE-2.0 - * - * Unless required by applicable law or agreed to in writing, software - * distributed under the License is distributed on an "AS IS" BASIS, - * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. - * See the License for the specific language governing permissions and - * limitations under the License. - */ - -#include - -#include - -namespace cucascade { -namespace memory { -namespace experimental { - -reservation_aware_resource_adaptor::reservation_aware_resource_adaptor( - any_device_resource primary_mr, std::int64_t limit) - : shared_base(::cuda::mr::make_shared_resource(std::move(primary_mr), limit)) -{ -} - -std::int64_t reservation_aware_resource_adaptor::limit() const noexcept { return get().limit(); } - -void reservation_aware_resource_adaptor::set_limit(std::int64_t limit) noexcept -{ - get().set_limit(limit); -} - -std::int64_t reservation_aware_resource_adaptor::current_allocated() const noexcept -{ - return get().current_allocated(); -} - -std::int64_t reservation_aware_resource_adaptor::total_reserved() const noexcept -{ - return get().total_reserved(); -} - -std::int64_t reservation_aware_resource_adaptor::available() const noexcept -{ - return get().available(); -} - -memory_record reservation_aware_resource_adaptor::get_main_record() const -{ - return get().get_main_record(); -} - -rmm::device_async_resource_ref reservation_aware_resource_adaptor::get_upstream_resource() - const noexcept -{ - return rmm::device_async_resource_ref{ - const_cast(get().get_upstream_resource())}; -} - -memory_reservation::memory_reservation(reservation_aware_resource_adaptor const& adaptor, - std::size_t granted, - std::size_t overbooking) - : shared_base{::cuda::mr::make_shared_resource( - adaptor, detail::safe_cast(granted), overbooking)} -{ -} - -memory_reservation reservation_aware_resource_adaptor::reserve(std::size_t size, - allow_overbooking overbooking_policy) -{ - auto const [granted, overbooking] = - get().reserve(size, overbooking_policy == allow_overbooking::YES); - return memory_reservation{*this, granted, overbooking}; -} - -} // namespace experimental -} // namespace memory -} // namespace cucascade diff --git a/test/memory/test_experimental_reservation_aware_resource_adaptor.cpp b/test/memory/test_experimental_reservation_aware_resource_adaptor.cpp index e9f34a4..d1d2a93 100644 --- a/test/memory/test_experimental_reservation_aware_resource_adaptor.cpp +++ b/test/memory/test_experimental_reservation_aware_resource_adaptor.cpp @@ -30,6 +30,7 @@ #include #include #include +#include #include #include @@ -42,12 +43,13 @@ #include #include -using cucascade::memory::experimental::allow_overbooking; -using cucascade::memory::experimental::memory_reservation; -using cucascade::memory::experimental::reservation_aware_resource_adaptor; +using namespace cucascade::memory::experimental; namespace { +/// The concrete state a reservation's erased handle wraps, for `resource_cast` round trips. +using device_reservation_handle = detail::reservation_handle; + bool has_cuda_device() { int device_count = 0; @@ -69,15 +71,23 @@ TEST_CASE("Reserve moves bytes from available to reserved", "[experimental_reser { if (!has_cuda_device()) { return; } - reservation_aware_resource_adaptor adaptor{ - ::cuda::mr::any_resource<::cuda::mr::device_accessible>{rmm::mr::cuda_memory_resource{}}, - limit}; + device_adaptor adaptor{any_device_resource{rmm::mr::cuda_memory_resource{}}, limit}; auto check_any_resouce_conversion = [&](auto& mr_like) { + using mr_like_t = std::remove_cvref_t; cuda::mr::any_resource any_device = mr_like; - cuda::mr::any_resource<> any_adaptor = mr_like; + cuda::mr::any_resource<> any_erased = mr_like; CHECK(any_device == mr_like); - CHECK(any_adaptor == mr_like); + CHECK(any_erased == mr_like); + + auto res_cast = [](auto* _res_ptr) { + auto casted_ptr = cuda::mr::resource_cast(_res_ptr); + CHECK(casted_ptr != nullptr); + return casted_ptr; + }; + + CHECK(mr_like == *res_cast(&any_device)); + CHECK(mr_like == *res_cast(&any_erased)); }; check_any_resouce_conversion(adaptor); @@ -86,7 +96,23 @@ TEST_CASE("Reserve moves bytes from available to reserved", "[experimental_reser REQUIRE_NOTHROW(std::ignore = adaptor.reserve(0, allow_overbooking::NO)); auto res = adaptor.reserve(1024, allow_overbooking::NO); - check_any_resouce_conversion(res); + + CHECK(res.accessibility() == reservation_accessibility::DEVICE); + CHECK(res.is_device_accessible()); + CHECK_FALSE(res.is_host_accessible()); + CHECK_THROWS_AS(res.as_host(), cucascade::logic_error); + CHECK_THROWS_AS(res.as_host_device(), cucascade::logic_error); + + // The projection owns a reference to the same shared state, recoverable by name. + auto projected = res.as_device(); + cuda::mr::any_resource<> erased = projected; + CHECK(erased == projected); + auto* handle = cuda::mr::resource_cast(&projected); + REQUIRE(handle != nullptr); + CHECK((*handle)->balance() == 1024); + + auto copy = res; + CHECK(copy == res); CHECK(res.overbooking() == 0); CHECK(res.balance() == 1024); @@ -99,20 +125,18 @@ TEST_CASE("Allocating keeps available unchanged", "[experimental_reservation_awa if (!has_cuda_device()) { return; } rmm::cuda_stream_view stream{rmm::cuda_stream_default}; - reservation_aware_resource_adaptor adaptor{ - ::cuda::mr::any_resource<::cuda::mr::device_accessible>{rmm::mr::cuda_memory_resource{}}, - limit}; + device_adaptor adaptor{any_device_resource{rmm::mr::cuda_memory_resource{}}, limit}; auto res = adaptor.reserve(1024, allow_overbooking::NO); { - rmm::device_buffer buf1{256, stream, res}; + rmm::device_buffer buf1{256, stream, res.as_device()}; CHECK(res.balance() == 768); CHECK(adaptor.total_reserved() == 768); CHECK(adaptor.current_allocated() == 256); CHECK(adaptor.available() == limit - 1024); - rmm::device_buffer buf2{512, stream, res}; + rmm::device_buffer buf2{512, stream, res.as_device()}; CHECK(res.balance() == 256); CHECK(adaptor.total_reserved() == 256); CHECK(adaptor.current_allocated() == 768); @@ -130,12 +154,10 @@ TEST_CASE("Exceeding the grant throws", "[experimental_reservation_aware][gpu]") if (!has_cuda_device()) { return; } rmm::cuda_stream_view stream{rmm::cuda_stream_default}; - reservation_aware_resource_adaptor adaptor{ - ::cuda::mr::any_resource<::cuda::mr::device_accessible>{rmm::mr::cuda_memory_resource{}}, - limit}; + device_adaptor adaptor{any_device_resource{rmm::mr::cuda_memory_resource{}}, limit}; auto res = adaptor.reserve(1024, allow_overbooking::NO); - REQUIRE_THROWS_AS((rmm::device_buffer{2048, stream, res}), rmm::out_of_memory); + REQUIRE_THROWS_AS((rmm::device_buffer{2048, stream, res.as_device()}), rmm::out_of_memory); CHECK(res.balance() == 1024); CHECK(adaptor.current_allocated() == 0); stream.synchronize(); @@ -146,14 +168,12 @@ TEST_CASE("Zero-sized reservation throws on first byte", "[experimental_reservat if (!has_cuda_device()) { return; } rmm::cuda_stream_view stream{rmm::cuda_stream_default}; - reservation_aware_resource_adaptor adaptor{ - ::cuda::mr::any_resource<::cuda::mr::device_accessible>{rmm::mr::cuda_memory_resource{}}, - limit}; + device_adaptor adaptor{any_device_resource{rmm::mr::cuda_memory_resource{}}, limit}; auto res = adaptor.reserve(static_cast(2 * limit), allow_overbooking::NO); CHECK(res.balance() == 0); CHECK(res.overbooking() == static_cast(limit)); - REQUIRE_THROWS_AS((rmm::device_buffer{1, stream, res}), rmm::out_of_memory); + REQUIRE_THROWS_AS((rmm::device_buffer{1, stream, res.as_device()}), rmm::out_of_memory); stream.synchronize(); } @@ -161,9 +181,7 @@ TEST_CASE("Overbooking is granted when allowed", "[experimental_reservation_awar { if (!has_cuda_device()) { return; } - reservation_aware_resource_adaptor adaptor{ - ::cuda::mr::any_resource<::cuda::mr::device_accessible>{rmm::mr::cuda_memory_resource{}}, - limit}; + device_adaptor adaptor{any_device_resource{rmm::mr::cuda_memory_resource{}}, limit}; auto res = adaptor.reserve(static_cast(2 * limit), allow_overbooking::YES); CHECK(res.balance() == static_cast(2 * limit)); @@ -171,13 +189,57 @@ TEST_CASE("Overbooking is granted when allowed", "[experimental_reservation_awar CHECK(adaptor.available() == -limit); } +TEST_CASE("Host reservations project to host only", "[experimental_reservation_aware][gpu]") +{ + if (!has_cuda_device()) { return; } + + host_adaptor adaptor{any_host_resource{rmm::mr::pinned_host_memory_resource{}}, limit}; + + auto res = adaptor.reserve(1024, allow_overbooking::NO); + CHECK(res.accessibility() == reservation_accessibility::HOST); + CHECK(res.is_host_accessible()); + CHECK_FALSE(res.is_device_accessible()); + CHECK_THROWS_AS(res.as_device(), cucascade::logic_error); + CHECK_THROWS_AS(res.as_host_device(), cucascade::logic_error); + + auto mr = res.as_host(); + auto* data = mr.allocate_sync(256, 256); + CHECK(res.balance() == 768); + CHECK(adaptor.current_allocated() == 256); + mr.deallocate_sync(data, 256, 256); + CHECK(res.balance() == 1024); + CHECK(adaptor.current_allocated() == 0); +} + +TEST_CASE("Host-device reservations project three ways", "[experimental_reservation_aware][gpu]") +{ + if (!has_cuda_device()) { return; } + + rmm::cuda_stream_view stream{rmm::cuda_stream_default}; + host_device_adaptor adaptor{any_host_device_resource{rmm::mr::pinned_host_memory_resource{}}, + limit}; + + auto res = adaptor.reserve(1024, allow_overbooking::NO); + CHECK(res.accessibility() == reservation_accessibility::HOST_DEVICE); + CHECK(res.is_host_accessible()); + CHECK(res.is_device_accessible()); + CHECK_NOTHROW(std::ignore = res.as_host()); + CHECK_NOTHROW(std::ignore = res.as_host_device()); + + { + rmm::device_buffer buf{256, stream, res.as_device()}; + CHECK(res.balance() == 768); + CHECK(adaptor.current_allocated() == 256); + } + CHECK(res.balance() == 1024); + stream.synchronize(); +} + TEST_CASE("Destruction refunds the unused balance", "[experimental_reservation_aware][gpu]") { if (!has_cuda_device()) { return; } - reservation_aware_resource_adaptor adaptor{ - ::cuda::mr::any_resource<::cuda::mr::device_accessible>{rmm::mr::cuda_memory_resource{}}, - limit}; + device_adaptor adaptor{any_device_resource{rmm::mr::cuda_memory_resource{}}, limit}; { auto res = adaptor.reserve(1024, allow_overbooking::NO); @@ -192,20 +254,19 @@ TEST_CASE("Buffer outlives the reserving scope", "[experimental_reservation_awar if (!has_cuda_device()) { return; } rmm::cuda_stream_view stream{rmm::cuda_stream_default}; - reservation_aware_resource_adaptor adaptor{ - ::cuda::mr::any_resource<::cuda::mr::device_accessible>{rmm::mr::cuda_memory_resource{}}, - limit}; + device_adaptor adaptor{any_device_resource{rmm::mr::cuda_memory_resource{}}, limit}; { auto buf = [&] { auto res = adaptor.reserve(1024, allow_overbooking::NO); - return rmm::device_buffer{512, stream, res}; + return rmm::device_buffer{512, stream, res.as_device()}; }(); - auto mr = buf.memory_resource(); - auto* reservation = ::cuda::mr::resource_cast(&mr); - REQUIRE(reservation != nullptr); - CHECK(reservation->balance() == 512); + // The buffer's own handle is now the only reference keeping the reservation alive. + auto mr = buf.memory_resource(); + auto* handle = ::cuda::mr::resource_cast(&mr); + REQUIRE(handle != nullptr); + CHECK((*handle)->balance() == 512); CHECK(adaptor.current_allocated() == 512); CHECK(adaptor.total_reserved() == 512); @@ -225,13 +286,11 @@ TEST_CASE("Main memory record tracks allocations", "[experimental_reservation_aw if (!has_cuda_device()) { return; } rmm::cuda_stream_view stream{rmm::cuda_stream_default}; - reservation_aware_resource_adaptor adaptor{ - ::cuda::mr::any_resource<::cuda::mr::device_accessible>{rmm::mr::cuda_memory_resource{}}, - limit}; + device_adaptor adaptor{any_device_resource{rmm::mr::cuda_memory_resource{}}, limit}; auto res = adaptor.reserve(1024, allow_overbooking::NO); { - rmm::device_buffer buf{256, stream, res}; + rmm::device_buffer buf{256, stream, res.as_device()}; auto record = adaptor.get_main_record(); CHECK(record.current == 256); CHECK(record.total == 256); @@ -259,9 +318,7 @@ TEST_CASE("Concurrent allocations share one reservation", "[experimental_reserva constexpr std::size_t num_threads = 2; constexpr std::size_t grant = num_buffers * max_buffer_size; - reservation_aware_resource_adaptor adaptor{ - ::cuda::mr::any_resource<::cuda::mr::device_accessible>{rmm::mr::cuda_memory_resource{}}, - limit}; + device_adaptor adaptor{any_device_resource{rmm::mr::cuda_memory_resource{}}, limit}; std::mt19937 rng{42}; std::uniform_int_distribution dist{0, max_buffer_size}; @@ -280,7 +337,7 @@ TEST_CASE("Concurrent allocations share one reservation", "[experimental_reserva workers.push_back(std::async(std::launch::async, [&, tid] { for (std::size_t i = tid; i < num_buffers; i += num_threads) { auto alloc_stream = pool.get_stream(i % pool.get_pool_size()); - buffers[i] = rmm::device_buffer{sizes[i], alloc_stream, res}; + buffers[i] = rmm::device_buffer{sizes[i], alloc_stream, res.as_device()}; } })); } From 8110d1d0207bbac7209c1e6612496720a4ccedaf Mon Sep 17 00:00:00 2001 From: niranda perera Date: Wed, 5 Aug 2026 15:52:22 -0700 Subject: [PATCH 4/6] minor change Signed-off-by: niranda perera --- .../reservation_aware_resource_adaptor.hpp | 32 +++++++++---------- 1 file changed, 16 insertions(+), 16 deletions(-) diff --git a/include/cucascade/memory/experimental/reservation_aware_resource_adaptor.hpp b/include/cucascade/memory/experimental/reservation_aware_resource_adaptor.hpp index 9307f03..b8f3a44 100644 --- a/include/cucascade/memory/experimental/reservation_aware_resource_adaptor.hpp +++ b/include/cucascade/memory/experimental/reservation_aware_resource_adaptor.hpp @@ -204,27 +204,27 @@ class reservation_aware_resource_adaptor_impl * @return A pair of the number of bytes granted (either @p size or zero) and the * number of bytes by which the request overbooks the limit. * - * @note Rejections are best-effort under contention: concurrent requests each claim - * before checking, so requests that would fit individually may be rejected. + * @note The decision is made against a snapshot: `limit_` and `current_` are read + * separately from the commit, so a concurrent allocation can still push the total + * past the limit after a request is granted. */ [[nodiscard]] std::pair reserve(std::size_t size, bool allow_overbooking) { - auto const want = safe_cast(size); - std::int64_t const capacity = limit() - current_allocated(); - - // Claim the bytes up front and roll back if they didn't fit. While a claim is - // being rolled back the reserved total reads high, which makes a concurrent - // `available()` pessimistic, never optimistic. - auto const reserved = total_reserved_.add(want, std::memory_order_acq_rel) - want; - std::int64_t const headroom = capacity - (reserved + want); - if (headroom >= 0) { return {size, 0}; } - auto const overbooking = safe_cast(-headroom); - if (!allow_overbooking) { - total_reserved_.sub(want, std::memory_order_acq_rel); - return {0, overbooking}; + auto const want = safe_cast(size); + auto reserved = total_reserved_.load(std::memory_order_acquire); + + // Commit the claim only once it is known to fit, so a rejected request never + // writes to the counter and never inflates what a concurrent `available()` sees. + // A failed exchange re-reads the limit and the allocated total as well. + while (true) { + std::int64_t const headroom = limit() - current_allocated() - reserved - want; + if (headroom < 0 && !allow_overbooking) { return {0, safe_cast(-headroom)}; } + if (total_reserved_->compare_exchange_weak( + reserved, reserved + want, std::memory_order_acq_rel, std::memory_order_acquire)) { + return {size, headroom < 0 ? safe_cast(-headroom) : 0}; + } } - return {size, overbooking}; } void* allocate(::cuda::stream_ref stream, From 2726aef2625322e5dd285c70f1b80c52383ae77b Mon Sep 17 00:00:00 2001 From: niranda perera Date: Wed, 5 Aug 2026 17:21:45 -0700 Subject: [PATCH 5/6] minor fixes Signed-off-by: niranda perera --- .../memory/experimental/memory_reservation.hpp | 4 ++-- .../reservation_aware_resource_adaptor.hpp | 5 +++-- .../reservation_aware_resource_adaptor.cpp | 11 ++++++++++- 3 files changed, 15 insertions(+), 5 deletions(-) diff --git a/include/cucascade/memory/experimental/memory_reservation.hpp b/include/cucascade/memory/experimental/memory_reservation.hpp index fb74a00..22ff150 100644 --- a/include/cucascade/memory/experimental/memory_reservation.hpp +++ b/include/cucascade/memory/experimental/memory_reservation.hpp @@ -122,10 +122,9 @@ class memory_reservation_impl return this == std::addressof(other); } + /// @brief The hook `cuda::forward_property` uses to forward the adaptor's properties. [[nodiscard]] Adaptor const& upstream_resource() const noexcept { return adaptor_; } - [[nodiscard]] Adaptor const& adaptor() const noexcept { return adaptor_; } - private: void draw_down_res(std::int64_t bytes) { @@ -278,6 +277,7 @@ class memory_reservation { private: template + requires ::cuda::mr::resource friend class reservation_aware_resource_adaptor; /// @brief The shared state, one alternative per accessibility. diff --git a/include/cucascade/memory/experimental/reservation_aware_resource_adaptor.hpp b/include/cucascade/memory/experimental/reservation_aware_resource_adaptor.hpp index b8f3a44..d95d7ed 100644 --- a/include/cucascade/memory/experimental/reservation_aware_resource_adaptor.hpp +++ b/include/cucascade/memory/experimental/reservation_aware_resource_adaptor.hpp @@ -43,6 +43,7 @@ namespace memory { namespace experimental { template + requires ::cuda::mr::resource class reservation_aware_resource_adaptor; class memory_reservation; @@ -157,10 +158,9 @@ class reservation_aware_resource_adaptor_impl return this == std::addressof(other); } + /// @brief The hook `cuda::forward_property` uses to forward the upstream's properties. [[nodiscard]] Upstream const& upstream_resource() const noexcept { return upstream_mr_; } - [[nodiscard]] Upstream const& get_upstream_resource() const noexcept { return upstream_mr_; } - [[nodiscard]] std::int64_t limit() const noexcept { return limit_.load(std::memory_order_acquire); @@ -335,6 +335,7 @@ class reservation_aware_resource_adaptor_impl * leaving `available()` unchanged; that is what makes a reservation a promise. */ template + requires ::cuda::mr::resource class reservation_aware_resource_adaptor : public ::cuda::mr::shared_resource> { public: diff --git a/src/memory/experimental/reservation_aware_resource_adaptor.cpp b/src/memory/experimental/reservation_aware_resource_adaptor.cpp index 58982ab..74e5d40 100644 --- a/src/memory/experimental/reservation_aware_resource_adaptor.cpp +++ b/src/memory/experimental/reservation_aware_resource_adaptor.cpp @@ -25,6 +25,7 @@ namespace memory { namespace experimental { template + requires ::cuda::mr::resource reservation_aware_resource_adaptor::reservation_aware_resource_adaptor( Upstream primary_mr, std::int64_t limit) : shared_base(::cuda::mr::make_shared_resource(std::move(primary_mr), limit)) @@ -32,48 +33,56 @@ reservation_aware_resource_adaptor::reservation_aware_resource_adaptor } template + requires ::cuda::mr::resource std::int64_t reservation_aware_resource_adaptor::limit() const noexcept { return this->get().limit(); } template + requires ::cuda::mr::resource void reservation_aware_resource_adaptor::set_limit(std::int64_t limit) noexcept { this->get().set_limit(limit); } template + requires ::cuda::mr::resource std::int64_t reservation_aware_resource_adaptor::current_allocated() const noexcept { return this->get().current_allocated(); } template + requires ::cuda::mr::resource std::int64_t reservation_aware_resource_adaptor::total_reserved() const noexcept { return this->get().total_reserved(); } template + requires ::cuda::mr::resource std::int64_t reservation_aware_resource_adaptor::available() const noexcept { return this->get().available(); } template + requires ::cuda::mr::resource memory_record reservation_aware_resource_adaptor::get_main_record() const { return this->get().get_main_record(); } template + requires ::cuda::mr::resource Upstream const& reservation_aware_resource_adaptor::get_upstream_resource() const noexcept { - return this->get().get_upstream_resource(); + return this->get().upstream_resource(); } template + requires ::cuda::mr::resource memory_reservation reservation_aware_resource_adaptor::reserve( std::size_t size, allow_overbooking overbooking_policy) { From 6248569a7bb3bf9f5df0e7b0057fabd845d12b43 Mon Sep 17 00:00:00 2001 From: niranda perera Date: Wed, 5 Aug 2026 18:52:04 -0700 Subject: [PATCH 6/6] adding reserve_soft Signed-off-by: niranda perera --- .../experimental/memory_reservation.hpp | 67 ++++++++++++++--- .../reservation_aware_resource_adaptor.hpp | 55 +++++++++++++- .../experimental/memory_reservation.cpp | 10 ++- .../reservation_aware_resource_adaptor.cpp | 16 +++- ...tal_reservation_aware_resource_adaptor.cpp | 75 ++++++++++++++++++- 5 files changed, 201 insertions(+), 22 deletions(-) diff --git a/include/cucascade/memory/experimental/memory_reservation.hpp b/include/cucascade/memory/experimental/memory_reservation.hpp index 22ff150..e6bf389 100644 --- a/include/cucascade/memory/experimental/memory_reservation.hpp +++ b/include/cucascade/memory/experimental/memory_reservation.hpp @@ -39,6 +39,13 @@ namespace detail { * Allocating moves bytes from the adaptor's reserved counter to its allocated counter; * the unspent balance is refunded only when the last reference dies. * + * The reservation's claim on `reservation_aware_resource_adaptor::total_reserved()` is + * `reserved_part(balance())`, never the raw balance, so a soft reservation that has + * overdrawn claims nothing rather than crediting back memory it is still using. Every + * balance transition adjusts the counter by the change in that quantity, which keeps the + * claim exact across allocation, partial release, and full recovery, and unwinds it to + * zero on destruction. + * * @tparam Adaptor One of `device_adaptor`, `host_adaptor`, or `host_device_adaptor`. Its * properties are forwarded to the reservation via `cuda::forward_property`, so a * reservation advertises whatever the granting adaptor advertises (e.g. @@ -50,14 +57,21 @@ template class memory_reservation_impl : public ::cuda::forward_property, Adaptor> { public: - memory_reservation_impl(Adaptor adaptor, std::int64_t grant, std::size_t overbooking) - : adaptor_{std::move(adaptor)}, grant_{grant}, overbooking_{overbooking}, balance_{grant} + memory_reservation_impl(Adaptor adaptor, + std::int64_t grant, + std::size_t overbooking, + grant_enforcement enforcement) + : adaptor_{std::move(adaptor)}, + grant_{grant}, + overbooking_{overbooking}, + enforcement_{enforcement}, + balance_{grant} { } ~memory_reservation_impl() { - adaptor_->total_reserved_.sub(balance(), std::memory_order_acq_rel); + adaptor_->total_reserved_.sub(reserved_part(balance()), std::memory_order_acq_rel); } memory_reservation_impl(memory_reservation_impl const&) = delete; @@ -69,6 +83,8 @@ class memory_reservation_impl [[nodiscard]] std::size_t overbooking() const noexcept { return overbooking_; } + [[nodiscard]] bool is_soft() const noexcept { return enforcement_ == grant_enforcement::SOFT; } + [[nodiscard]] std::int64_t balance() const noexcept { return balance_.load(std::memory_order_acquire); @@ -79,15 +95,16 @@ class memory_reservation_impl std::size_t alignment = rmm::CUDA_ALLOCATION_ALIGNMENT) { auto const amount = safe_cast(bytes); - draw_down_res(amount); - void* ptr = nullptr; + auto const before = draw_down_res(amount); + void* ptr = nullptr; try { ptr = adaptor_->allocate(stream, bytes, alignment); } catch (...) { balance_.fetch_add(amount, std::memory_order_acq_rel); throw; } - adaptor_->total_reserved_.sub(amount, std::memory_order_acq_rel); + adaptor_->total_reserved_.sub(reserved_part(before) - reserved_part(before - amount), + std::memory_order_acq_rel); return ptr; } @@ -97,8 +114,9 @@ class memory_reservation_impl std::size_t alignment = rmm::CUDA_ALLOCATION_ALIGNMENT) noexcept { auto const amount = safe_cast(bytes); - balance_.fetch_add(amount, std::memory_order_acq_rel); - adaptor_->total_reserved_.add(amount, std::memory_order_acq_rel); + auto const before = balance_.fetch_add(amount, std::memory_order_acq_rel); + adaptor_->total_reserved_.add(reserved_part(before + amount) - reserved_part(before), + std::memory_order_acq_rel); adaptor_->deallocate(stream, ptr, bytes, alignment); } @@ -126,8 +144,20 @@ class memory_reservation_impl [[nodiscard]] Adaptor const& upstream_resource() const noexcept { return adaptor_; } private: - void draw_down_res(std::int64_t bytes) + /// @brief The part of a balance that is still reserved-but-unspent. An overdraft is + /// funded from outside the grant, so it contributes nothing to the adaptor's reserve. + static constexpr std::int64_t reserved_part(std::int64_t balance) noexcept + { + return balance > 0 ? balance : 0; + } + + /// @return The balance before the draw-down. + std::int64_t draw_down_res(std::int64_t bytes) { + // Nothing to enforce, so skip the compare-exchange entirely and let the balance + // go negative; the overdraft still shows up in the adaptor's `current_allocated()`. + if (is_soft()) { return balance_.fetch_sub(bytes, std::memory_order_acq_rel); } + auto balance = balance_.load(std::memory_order_relaxed); do { if (bytes > balance) { @@ -138,11 +168,13 @@ class memory_reservation_impl } } while (!balance_.compare_exchange_weak( balance, balance - bytes, std::memory_order_acq_rel, std::memory_order_relaxed)); + return balance; } Adaptor adaptor_; std::int64_t const grant_; std::size_t const overbooking_; + grant_enforcement const enforcement_; std::atomic balance_; }; @@ -162,6 +194,11 @@ using reservation_handle = ::cuda::mr::shared_resource @@ -389,6 +400,28 @@ class reservation_aware_resource_adaptor */ [[nodiscard]] memory_reservation reserve(std::size_t size, allow_overbooking overbooking_policy); + /** + * @brief Reserve an amount of memory without capping allocations at the grant. + * + * Identical to `reserve()` in how the grant is sized and accounted, but allocations + * through the returned reservation are never refused for exceeding it. Going past the + * grant drives `memory_reservation::balance()` negative by the overdraft. The overdrawn + * bytes count only in `current_allocated()` and contribute nothing to + * `total_reserved()`, so `available()` treats the overdraft as consumed rather than as + * free space for as long as it lasts. + * + * Nothing throttles an overdraft. It is charged against the adaptor's limit but not + * bounded by it, so `available()` can go negative and stay there until the memory is + * released. Use this when the total is not known up front and a hard failure + * mid-pipeline is worse than temporarily exceeding the budget. + * + * @param size The number of bytes to reserve. + * @param overbooking_policy Whether overbooking is allowed. + * @return The reservation, which reports `memory_reservation::is_soft() == true`. + */ + [[nodiscard]] memory_reservation reserve_soft(std::size_t size, + allow_overbooking overbooking_policy); + /** * @brief Get the memory limit. * @@ -411,10 +444,19 @@ class reservation_aware_resource_adaptor [[nodiscard]] std::int64_t current_allocated() const noexcept; /** - * @brief Get the memory currently held by live reservations. + * @brief Get the memory promised to live reservations but not yet allocated. * * Excludes reserved bytes that have already been allocated; those are reported by - * `current_allocated()` instead. + * `current_allocated()` instead. A soft reservation that has overdrawn its grant + * contributes zero rather than a negative amount, so an overdraft never reads back as + * free capacity. + * + * @note Accurate once the accounting settles. Each allocation updates the reservation's + * balance and this counter as two separate atomic operations, and concurrent + * allocations apply their updates in an order unrelated to the order they observed + * those balances. A read taken while updates are in flight can therefore land low, even + * below zero, by up to the size of the concurrent allocations. It resolves as they + * complete. * * @return Total number of reserved bytes. */ @@ -423,8 +465,13 @@ class reservation_aware_resource_adaptor /** * @brief Get the memory available for new reservations. * - * Computed as `limit() - current_allocated() - total_reserved()`. May be negative - * when reservations have overbooked the limit. + * Computed as `limit() - current_allocated() - total_reserved()`. Negative when + * reservations have overbooked the limit or a soft reservation has overdrawn its grant. + * + * @note A best-effort snapshot, not an atomic one: the three counters are read + * independently and each carries the caveat on `total_reserved()`. Treat the result as + * a hint that a concurrent allocation may already have invalidated, which is why + * `reserve()` re-reads it rather than trusting a value handed in from outside. * * @return The available memory in bytes. */ diff --git a/src/memory/experimental/memory_reservation.cpp b/src/memory/experimental/memory_reservation.cpp index 24f0f86..455b2e3 100644 --- a/src/memory/experimental/memory_reservation.cpp +++ b/src/memory/experimental/memory_reservation.cpp @@ -133,10 +133,14 @@ std::size_t memory_reservation::grant() const noexcept [](auto const& handle) { return detail::safe_cast(handle->grant()); }, handle_); } -std::size_t memory_reservation::balance() const noexcept +std::int64_t memory_reservation::balance() const noexcept { - return std::visit( - [](auto const& handle) { return detail::safe_cast(handle->balance()); }, handle_); + return std::visit([](auto const& handle) { return handle->balance(); }, handle_); +} + +bool memory_reservation::is_soft() const noexcept +{ + return std::visit([](auto const& handle) { return handle->is_soft(); }, handle_); } std::size_t memory_reservation::overbooking() const noexcept diff --git a/src/memory/experimental/reservation_aware_resource_adaptor.cpp b/src/memory/experimental/reservation_aware_resource_adaptor.cpp index 74e5d40..60f5c16 100644 --- a/src/memory/experimental/reservation_aware_resource_adaptor.cpp +++ b/src/memory/experimental/reservation_aware_resource_adaptor.cpp @@ -92,7 +92,21 @@ memory_reservation reservation_aware_resource_adaptor::reserve( using impl_t = detail::memory_reservation_impl>; return memory_reservation{ memory_reservation::handle_variant{::cuda::mr::make_shared_resource( - *this, detail::safe_cast(granted), overbooking)}}; + *this, detail::safe_cast(granted), overbooking, grant_enforcement::STRICT)}}; +} + +template + requires ::cuda::mr::resource +memory_reservation reservation_aware_resource_adaptor::reserve_soft( + std::size_t size, allow_overbooking overbooking_policy) +{ + auto const [granted, overbooking] = + this->get().reserve(size, overbooking_policy == allow_overbooking::YES); + + using impl_t = detail::memory_reservation_impl>; + return memory_reservation{ + memory_reservation::handle_variant{::cuda::mr::make_shared_resource( + *this, detail::safe_cast(granted), overbooking, grant_enforcement::SOFT)}}; } template class reservation_aware_resource_adaptor; diff --git a/test/memory/test_experimental_reservation_aware_resource_adaptor.cpp b/test/memory/test_experimental_reservation_aware_resource_adaptor.cpp index d1d2a93..a01996f 100644 --- a/test/memory/test_experimental_reservation_aware_resource_adaptor.cpp +++ b/test/memory/test_experimental_reservation_aware_resource_adaptor.cpp @@ -163,6 +163,73 @@ TEST_CASE("Exceeding the grant throws", "[experimental_reservation_aware][gpu]") stream.synchronize(); } +TEST_CASE("Soft reservations allow exceeding the grant", "[experimental_reservation_aware][gpu]") +{ + if (!has_cuda_device()) { return; } + + rmm::cuda_stream_view stream{rmm::cuda_stream_default}; + device_adaptor adaptor{any_device_resource{rmm::mr::cuda_memory_resource{}}, limit}; + + auto res = adaptor.reserve_soft(1024, allow_overbooking::NO); + CHECK(res.is_soft()); + CHECK(res.balance() == 1024); + + { + // Overdrawing is permitted and shows up as a negative balance. + rmm::device_buffer buf{3072, stream, res.as_device()}; + CHECK(res.balance() == -2048); + + // The overdraft shows up as consumed memory, not as returned reserve. + CHECK(adaptor.current_allocated() == 3072); + CHECK(adaptor.total_reserved() == 0); + CHECK(adaptor.available() == limit - 3072); + } + + CHECK(res.balance() == 1024); + CHECK(adaptor.current_allocated() == 0); + CHECK(adaptor.total_reserved() == 1024); + stream.synchronize(); +} + +TEST_CASE("Overdrawn soft reservation outlived by its buffer", + "[experimental_reservation_aware][gpu]") +{ + if (!has_cuda_device()) { return; } + + rmm::cuda_stream_view stream{rmm::cuda_stream_default}; + device_adaptor adaptor{any_device_resource{rmm::mr::cuda_memory_resource{}}, limit}; + + { + rmm::device_buffer buf = [&] { + auto res = adaptor.reserve_soft(1024, allow_overbooking::NO); + return rmm::device_buffer{4096, stream, res.as_device()}; + }(); + // Only the reservation handle is gone; the buffer still holds the shared state, so + // the refund has not run yet. The grant is fully drawn, hence a zero reserve. + CHECK(adaptor.total_reserved() == 0); + CHECK(adaptor.current_allocated() == 4096); + } + + CHECK(adaptor.total_reserved() == 0); + CHECK(adaptor.current_allocated() == 0); + CHECK(adaptor.available() == limit); + stream.synchronize(); +} + +TEST_CASE("Strict reservations remain capped", "[experimental_reservation_aware][gpu]") +{ + if (!has_cuda_device()) { return; } + + rmm::cuda_stream_view stream{rmm::cuda_stream_default}; + device_adaptor adaptor{any_device_resource{rmm::mr::cuda_memory_resource{}}, limit}; + + auto res = adaptor.reserve(1024, allow_overbooking::NO); + CHECK_FALSE(res.is_soft()); + REQUIRE_THROWS_AS((rmm::device_buffer{3072, stream, res.as_device()}), rmm::out_of_memory); + CHECK(res.balance() == 1024); + stream.synchronize(); +} + TEST_CASE("Zero-sized reservation throws on first byte", "[experimental_reservation_aware][gpu]") { if (!has_cuda_device()) { return; } @@ -184,7 +251,7 @@ TEST_CASE("Overbooking is granted when allowed", "[experimental_reservation_awar device_adaptor adaptor{any_device_resource{rmm::mr::cuda_memory_resource{}}, limit}; auto res = adaptor.reserve(static_cast(2 * limit), allow_overbooking::YES); - CHECK(res.balance() == static_cast(2 * limit)); + CHECK(res.balance() == 2 * limit); CHECK(res.overbooking() == static_cast(limit)); CHECK(adaptor.available() == -limit); } @@ -327,7 +394,7 @@ TEST_CASE("Concurrent allocations share one reservation", "[experimental_reserva auto const total = std::accumulate(sizes.begin(), sizes.end(), std::size_t{0}); auto res = adaptor.reserve(grant, allow_overbooking::NO); - REQUIRE(res.balance() == grant); + REQUIRE(res.balance() == static_cast(grant)); rmm::cuda_stream_pool pool{4, rmm::cuda_stream::flags::non_blocking}; std::vector buffers(num_buffers); @@ -345,13 +412,13 @@ TEST_CASE("Concurrent allocations share one reservation", "[experimental_reserva REQUIRE_NOTHROW(worker.get()); } - CHECK(res.balance() == grant - total); + CHECK(res.balance() == static_cast(grant - total)); CHECK(adaptor.total_reserved() == static_cast(grant - total)); CHECK(adaptor.current_allocated() == static_cast(total)); CHECK(adaptor.available() == limit - static_cast(grant)); buffers.clear(); - CHECK(res.balance() == grant); + CHECK(res.balance() == static_cast(grant)); CHECK(adaptor.current_allocated() == 0); synchronize_pool(pool);