From a5d8df2e1b2e573e9bd37b155b6213f3b88732a0 Mon Sep 17 00:00:00 2001 From: niranda perera Date: Wed, 29 Jul 2026 17:52:57 -0700 Subject: [PATCH] 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); +}