diff --git a/include/cucascade/memory/experimental/memory_reservation.hpp b/include/cucascade/memory/experimental/memory_reservation.hpp new file mode 100644 index 0000000..35bf29f --- /dev/null +++ b/include/cucascade/memory/experimental/memory_reservation.hpp @@ -0,0 +1,350 @@ +/* + * 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 + +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. + * std::shared_ptr policy); + + * 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. + * `cuda::mr::device_accessible`). That forwarding is what lets `memory_reservation` + * project this impl into an erased resource of the matching accessibility. + */ +template + requires reservation_adaptor +class memory_reservation_impl + : public reservation_control, + public ::cuda::forward_property, Adaptor> { + public: + memory_reservation_impl(Adaptor adaptor, + std::int64_t grant, + std::size_t overbooking, + std::shared_ptr policy) + : adaptor_{std::move(adaptor)}, + grant_{grant}, + overbooking_{overbooking}, + policy_{std::move(policy)}, + balance_{grant} + { + if (!policy_) { CUCASCADE_FAIL("over_reservation_policy must not be null"); } + } + + ~memory_reservation_impl() + { + adaptor_->total_reserved_.sub(reserved_part(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 override { return grant_; } + + [[nodiscard]] std::int64_t balance() const noexcept override + { + return balance_.load(std::memory_order_acquire); + } + + [[nodiscard]] std::size_t overbooking() const noexcept override { return overbooking_; } + + void* allocate(::cuda::stream_ref stream, + std::size_t bytes, + std::size_t alignment = rmm::CUDA_ALLOCATION_ALIGNMENT) + { + auto const amount = safe_cast(bytes); + auto const before = draw_down_res(amount); + auto const after = before - amount; + auto const claim_delta = calc_delta(before, after); + + adaptor_->total_reserved_.sub(claim_delta, std::memory_order_acq_rel); + + void* ptr = nullptr; + try { + ptr = adaptor_->allocate(stream, bytes, alignment); + } catch (...) { + auto const rollback_before = balance_.fetch_add(amount, std::memory_order_acq_rel); + auto const rollback_delta = calc_delta(rollback_before + amount, rollback_before); + adaptor_->total_reserved_.add(rollback_delta, std::memory_order_acq_rel); + throw; + } + 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); + auto const before = balance_.fetch_add(amount, std::memory_order_acq_rel); + adaptor_->total_reserved_.add(calc_delta(before + amount, before), 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); + } + + /// @brief The hook `cuda::forward_property` uses to forward the adaptor's properties. + [[nodiscard]] Adaptor const& upstream_resource() const noexcept { return adaptor_; } + + private: + /// @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 difference between two reserved portions. + static constexpr std::int64_t calc_delta(std::int64_t lhs, std::int64_t rhs) noexcept + { + return reserved_part(lhs) - reserved_part(rhs); + } + + /// @return The balance before the draw-down. + std::int64_t draw_down_res(std::int64_t bytes) + { + auto balance = balance_.load(std::memory_order_relaxed); + while (true) { + if (bytes > balance) { + policy_->handle_over_reservation(bytes, balance, *this); + return balance_.fetch_sub(bytes, std::memory_order_acq_rel); + } + + if (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_; + std::shared_ptr policy_; + 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, 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. 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. + * + * The default soft over-reservation policy permits allocations past the grant, reporting + * the overdraft as a negative `balance()`. Only the grant is backed by the adaptor's + * reserve; the overdraft is spent against the limit without a claim on it, so it shows up + * in `current_allocated()` alone. + * + * 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. + * + * @par Handing a reservation to cudf or RMM + * + * 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.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: + /** + * @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. + * + * @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; + + /** + * @brief The number of bytes originally granted. + * + * @return The granted size in bytes. + */ + [[nodiscard]] std::size_t grant() const noexcept; + + /** + * @brief The remaining unallocated size of the reservation. + * + * Negative on a soft reservation that has allocated past its grant, by the size of + * the overdraft. Never negative on a strict one, which refuses those allocations. + * + * @return The remaining size in bytes. + */ + [[nodiscard]] std::int64_t balance() const noexcept; + + /** + * @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; + + private: + template + requires ::cuda::mr::resource + 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 shared state. + * + * Private so that only `reservation_aware_resource_adaptor` can grant reservations. + * + * @param handle The shared state of the granted reservation. + */ + explicit memory_reservation(handle_variant handle) : handle_{std::move(handle)} {} + + handle_variant handle_; +}; + +} // namespace experimental +} // namespace memory +} // namespace cucascade diff --git a/include/cucascade/memory/experimental/over_reservation_policy.hpp b/include/cucascade/memory/experimental/over_reservation_policy.hpp new file mode 100644 index 0000000..21f8284 --- /dev/null +++ b/include/cucascade/memory/experimental/over_reservation_policy.hpp @@ -0,0 +1,108 @@ +/* + * 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 + +namespace cucascade { +namespace memory { +namespace experimental { + +/** + * @brief Narrow interface to reservation state and accounting operations. + * + * Implemented by the shared reservation state so policies can inspect the reservation + * without depending on the adaptor's accessibility. Future thread-safe grow/shrink + * operations belong on this interface because their implementations must coordinate + * reservation state with the granting adaptor's accounting. + */ +class reservation_control { + public: + virtual ~reservation_control() = default; + + [[nodiscard]] virtual std::int64_t grant() const noexcept = 0; + [[nodiscard]] virtual std::int64_t balance() const noexcept = 0; + [[nodiscard]] virtual std::size_t overbooking() const noexcept = 0; +}; + +/** + * @brief Pluggable policy invoked when an allocation would exceed the reservation balance. + * + * Called only after observing `requested_bytes > balance`. Returning normally authorizes + * the pending allocation; the allocation path then atomically subtracts the request from + * the current balance. Throwing `rmm::out_of_memory` rejects it. + */ +class over_reservation_policy { + public: + virtual ~over_reservation_policy(); + + /** + * @brief Handle an allocation that would exceed the reservation's remaining balance. + * + * @param requested_bytes Number of bytes being requested. + * @param observed_balance Remaining balance observed before invoking the policy. + * @param reservation The reservation whose balance would be exceeded. + * @throws rmm::out_of_memory if the policy decides to reject the allocation. + */ + virtual void handle_over_reservation(std::int64_t requested_bytes, + std::int64_t observed_balance, + reservation_control& reservation) const = 0; +}; + +/** + * @brief Hard policy — throws when the grant would be exceeded. + */ +class throw_on_over_reservation : public over_reservation_policy { + public: + throw_on_over_reservation(); + + void handle_over_reservation(std::int64_t requested_bytes, + std::int64_t observed_balance, + reservation_control& reservation) const final; +}; + +/** + * @brief Soft policy — allows the allocation to proceed past the grant. + */ +class ignore_on_over_reservation : public over_reservation_policy { + public: + ignore_on_over_reservation(); + + void handle_over_reservation(std::int64_t requested_bytes, + std::int64_t observed_balance, + reservation_control& reservation) const final; +}; + +/** + * @brief Shared stateless hard policy instance for `reserve()`. + */ +[[nodiscard]] std::shared_ptr const& +thow_on_over_reservation_instance(); + +/** + * @brief Shared stateless soft policy instance used by default reservations. + */ +[[nodiscard]] std::shared_ptr const& +ignore_on_over_reservation_instance(); + +} // 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 new file mode 100644 index 0000000..c6533c0 --- /dev/null +++ b/include/cucascade/memory/experimental/reservation_aware_resource_adaptor.hpp @@ -0,0 +1,480 @@ +/* + * 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 +#include +#include + +namespace cucascade { +namespace memory { +namespace experimental { + +template + requires ::cuda::mr::resource +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. + */ +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); + } +} + +/** + * @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 reservation_adaptor +class memory_reservation_impl; + +/** + * @brief Shared state of a reservation-aware resource adaptor. + * + * 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`. + */ +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. + * + * @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(Upstream 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); + } + + /// @brief The hook `cuda::forward_property` uses to forward the upstream's properties. + [[nodiscard]] Upstream const& 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 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); + 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}; + } + } + } + + 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(); + } + + private: + friend class memory_reservation_impl; + friend class memory_reservation_impl; + 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); + } + + Upstream 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}; +}; + +} // namespace detail + +/** + * @brief A memory resource adaptor that only allocates through reservations. + * + * 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 + * 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. + */ +template + requires ::cuda::mr::resource +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 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(Upstream 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 this->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. + * The per-reservation policy determines what happens when an allocation exceeds the + * remaining balance. The default soft policy permits the allocation; a custom policy + * may instead grow the reservation or throw. + * + * @param size The number of bytes to reserve. + * @param overbooking_policy Whether overbooking is allowed. + * @param policy Per-reservation policy invoked when an allocation exceeds the balance. + * Must not be null. + * @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, + std::shared_ptr policy = ignore_on_over_reservation_instance()); + + /** + * @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 promised to live reservations but not yet allocated. + * + * Excludes reserved bytes that have already been allocated; those are reported by + * `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. + */ + [[nodiscard]] std::int64_t total_reserved() const noexcept; + + /** + * @brief Get the memory available for new reservations. + * + * 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. + */ + [[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 erased upstream resource. + */ + [[nodiscard]] Upstream const& get_upstream_resource() const noexcept; + + private: + [[nodiscard]] memory_reservation make_reservation_impl( + std::size_t size, + allow_overbooking overbooking_policy, + std::shared_ptr policy); +}; + +} // namespace experimental +} // namespace memory +} // namespace cucascade + +#include diff --git a/src/memory/CMakeLists.txt b/src/memory/CMakeLists.txt index 03ba85e..db2630a 100644 --- a/src/memory/CMakeLists.txt +++ b/src/memory/CMakeLists.txt @@ -18,21 +18,25 @@ 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/memory_reservation.cpp + ${CMAKE_CURRENT_SOURCE_DIR}/experimental/over_reservation_policy.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/memory_reservation.cpp b/src/memory/experimental/memory_reservation.cpp new file mode 100644 index 0000000..7851012 --- /dev/null +++ b/src/memory/experimental/memory_reservation.cpp @@ -0,0 +1,148 @@ +/* + * 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::int64_t memory_reservation::balance() const noexcept +{ + return std::visit([](auto const& handle) { return 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/over_reservation_policy.cpp b/src/memory/experimental/over_reservation_policy.cpp new file mode 100644 index 0000000..01e202d --- /dev/null +++ b/src/memory/experimental/over_reservation_policy.cpp @@ -0,0 +1,83 @@ +/* + * 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 { + +over_reservation_policy::~over_reservation_policy() = default; + +throw_on_over_reservation::throw_on_over_reservation() = default; + +void throw_on_over_reservation::handle_over_reservation(std::int64_t requested_bytes, + std::int64_t observed_balance, + reservation_control& reservation) const +{ + CUCASCADE_FAIL("allocation of " + std::to_string(requested_bytes) + + " bytes exceeds reservation (grant: " + std::to_string(reservation.grant()) + + ", remaining: " + std::to_string(observed_balance) + ")", + rmm::out_of_memory); +} + +ignore_on_over_reservation::ignore_on_over_reservation() = default; + +void ignore_on_over_reservation::handle_over_reservation( + [[maybe_unused]] std::int64_t requested_bytes, + [[maybe_unused]] std::int64_t observed_balance, + [[maybe_unused]] reservation_control& reservation) const +{ + // Allow the pending allocation to proceed; the caller subtracts from balance. +} + +namespace { + +std::shared_ptr const make_hard_policy() +{ + static auto const policy = std::make_shared(); + return policy; +} + +std::shared_ptr const make_soft_policy() +{ + static auto const policy = std::make_shared(); + return policy; +} + +} // namespace + +std::shared_ptr const& thow_on_over_reservation_instance() +{ + static auto const policy = make_hard_policy(); + return policy; +} + +std::shared_ptr const& ignore_on_over_reservation_instance() +{ + static auto const policy = make_soft_policy(); + return policy; +} + +} // 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..926c8f7 --- /dev/null +++ b/src/memory/experimental/reservation_aware_resource_adaptor.cpp @@ -0,0 +1,123 @@ +/* + * 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 + 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)) +{ +} + +template + requires ::cuda::mr::resource +memory_reservation reservation_aware_resource_adaptor::make_reservation_impl( + std::size_t size, + allow_overbooking overbooking_policy, + std::shared_ptr 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, std::move(policy))}}; +} + +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().upstream_resource(); +} + +template + requires ::cuda::mr::resource +memory_reservation reservation_aware_resource_adaptor::reserve( + std::size_t size, + allow_overbooking overbooking_policy, + std::shared_ptr policy) +{ + return make_reservation_impl(size, overbooking_policy, std::move(policy)); +} + +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/test/CMakeLists.txt b/test/CMakeLists.txt index dd613aa..6f909bf 100644 --- a/test/CMakeLists.txt +++ b/test/CMakeLists.txt @@ -29,6 +29,7 @@ if(NOT CUCASCADE_TOPOLOGY_ONLY) add_executable( cucascade_tests # Memory tests + memory/test_experimental_reservation_aware_resource_adaptor.cpp memory/test_memory_reservation_manager.cpp memory/test_reservation_aware_resource_adaptor.cpp memory/test_reservation_manager_configurator.cpp 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..1fd6d0a --- /dev/null +++ b/test/memory/test_experimental_reservation_aware_resource_adaptor.cpp @@ -0,0 +1,557 @@ +/* + * 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 +#include +#include +#include +#include +#include +#include + +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; + 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; } + + 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_erased = mr_like; + CHECK(any_device == 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); + + REQUIRE(adaptor.available() == limit); + + REQUIRE_NOTHROW(std::ignore = adaptor.reserve(0, allow_overbooking::NO)); + + auto res = adaptor.reserve(1024, allow_overbooking::NO, thow_on_over_reservation_instance()); + + 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); + 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}; + 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.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.as_device()}; + 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}; + device_adaptor adaptor{any_device_resource{rmm::mr::cuda_memory_resource{}}, limit}; + + auto res = adaptor.reserve(1024, allow_overbooking::NO, thow_on_over_reservation_instance()); + 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(); +} + +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(1024, allow_overbooking::NO); + 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(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, thow_on_over_reservation_instance()); + 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; } + + 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(static_cast(2 * limit), + allow_overbooking::NO, + thow_on_over_reservation_instance()); + CHECK(res.balance() == 0); + CHECK(res.overbooking() == static_cast(limit)); + REQUIRE_THROWS_AS((rmm::device_buffer{1, stream, res.as_device()}), rmm::out_of_memory); + stream.synchronize(); +} + +TEST_CASE("Overbooking is granted when allowed", "[experimental_reservation_aware][gpu]") +{ + if (!has_cuda_device()) { return; } + + 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() == 2 * limit); + CHECK(res.overbooking() == static_cast(limit)); + 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; } + + device_adaptor adaptor{any_device_resource{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}; + device_adaptor adaptor{any_device_resource{rmm::mr::cuda_memory_resource{}}, limit}; + + { + auto buf = [&] { + auto res = adaptor.reserve(1024, allow_overbooking::NO, thow_on_over_reservation_instance()); + return rmm::device_buffer{512, stream, res.as_device()}; + }(); + + // 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); + 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}; + device_adaptor adaptor{any_device_resource{rmm::mr::cuda_memory_resource{}}, limit}; + + auto res = adaptor.reserve(1024, allow_overbooking::NO, thow_on_over_reservation_instance()); + { + rmm::device_buffer buf{256, stream, res.as_device()}; + 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; + + 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}; + 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() == static_cast(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.as_device()}; + } + })); + } + for (auto& worker : workers) { + REQUIRE_NOTHROW(worker.get()); + } + + 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() == static_cast(grant)); + CHECK(adaptor.current_allocated() == 0); + + synchronize_pool(pool); +} + +namespace { + +class throwing_device_resource { + public: + void* allocate(::cuda::stream_ref, std::size_t, std::size_t = rmm::CUDA_ALLOCATION_ALIGNMENT) + { + throw std::runtime_error("simulated upstream allocation failure"); + } + + void deallocate(::cuda::stream_ref, + void*, + std::size_t, + std::size_t = rmm::CUDA_ALLOCATION_ALIGNMENT) noexcept + { + } + + void* allocate_sync(std::size_t bytes, std::size_t alignment = rmm::CUDA_ALLOCATION_ALIGNMENT) + { + return allocate(::cuda::stream_ref{cudaStream_t{nullptr}}, bytes, alignment); + } + + void deallocate_sync(void* ptr, + std::size_t bytes, + std::size_t alignment = rmm::CUDA_ALLOCATION_ALIGNMENT) noexcept + { + deallocate(::cuda::stream_ref{cudaStream_t{nullptr}}, ptr, bytes, alignment); + } + + bool operator==(throwing_device_resource const&) const noexcept { return true; } + + friend void get_property(throwing_device_resource const&, ::cuda::mr::device_accessible) noexcept + { + } +}; + +class recording_over_reservation_policy : public over_reservation_policy { + public: + mutable std::mutex mutex_{}; + mutable std::int64_t last_requested_{0}; + mutable std::int64_t last_observed_{0}; + mutable std::int64_t last_grant_{0}; + mutable std::size_t call_count_{0}; + bool allow_{false}; + + void handle_over_reservation(std::int64_t requested_bytes, + std::int64_t observed_balance, + reservation_control& reservation) const override + { + std::lock_guard lock{mutex_}; + last_requested_ = requested_bytes; + last_observed_ = observed_balance; + last_grant_ = reservation.grant(); + ++call_count_; + if (!allow_) { throw rmm::out_of_memory{"recording policy rejects over-reservation"}; } + } +}; + +} // namespace + +TEST_CASE("Custom over-reservation policy receives reservation control", + "[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 policy = std::make_shared(); + policy->allow_ = true; + auto res = adaptor.reserve(1024, allow_overbooking::NO, policy); + + REQUIRE_NOTHROW((rmm::device_buffer{2048, stream, res.as_device()})); + CHECK(policy->call_count_ == 1); + CHECK(policy->last_requested_ == 2048); + CHECK(policy->last_observed_ == 1024); + CHECK(policy->last_grant_ == 1024); + stream.synchronize(); +} + +TEST_CASE("Concurrent insufficient draws on a hard reservation", + "[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, thow_on_over_reservation_instance()); + + std::vector> workers(2); + for (auto& worker : workers) { + worker = std::async(std::launch::async, [&] { + try { + return rmm::device_buffer{768, stream, res.as_device()}; + } catch (rmm::out_of_memory const&) { + return rmm::device_buffer{}; + } + }); + } + + std::vector buffers; + buffers.reserve(workers.size()); + std::ranges::transform( + workers, std::back_inserter(buffers), [](auto& worker) { return worker.get(); }); + + CHECK(std::ranges::count_if(buffers, [](auto const& buffer) { return buffer.size() > 0; }) == 1); + CHECK(res.balance() == 256); + CHECK(adaptor.current_allocated() == 768); + stream.synchronize(); +} + +TEST_CASE("Upstream allocation failure rolls back balance and reserve claim", + "[experimental_reservation_aware][gpu]") +{ + if (!has_cuda_device()) { return; } + + rmm::cuda_stream_view stream{rmm::cuda_stream_default}; + throwing_device_resource upstream{}; + device_adaptor adaptor{any_device_resource{upstream}, limit}; + auto res = adaptor.reserve(1024, allow_overbooking::NO); + + REQUIRE_THROWS_AS((rmm::device_buffer{256, stream, res.as_device()}), std::runtime_error); + CHECK(res.balance() == 1024); + CHECK(adaptor.total_reserved() == 1024); + CHECK(adaptor.current_allocated() == 0); + stream.synchronize(); +}