diff --git a/include/cucascade/memory/experimental/memory_reservation.hpp b/include/cucascade/memory/experimental/memory_reservation.hpp new file mode 100644 index 0000000..e6bf389 --- /dev/null +++ b/include/cucascade/memory/experimental/memory_reservation.hpp @@ -0,0 +1,349 @@ +/* + * 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. + * + * 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 ::cuda::forward_property, Adaptor> { + public: + memory_reservation_impl(Adaptor adaptor, + std::int64_t grant, + std::size_t overbooking, + grant_enforcement enforcement) + : adaptor_{std::move(adaptor)}, + grant_{grant}, + overbooking_{overbooking}, + enforcement_{enforcement}, + balance_{grant} + { + } + + ~memory_reservation_impl() + { + adaptor_->total_reserved_.sub(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 { return grant_; } + + [[nodiscard]] std::size_t overbooking() const noexcept { return overbooking_; } + + [[nodiscard]] bool is_soft() const noexcept { return enforcement_ == grant_enforcement::SOFT; } + + [[nodiscard]] std::int64_t balance() const noexcept + { + return balance_.load(std::memory_order_acquire); + } + + 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); + 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(reserved_part(before) - reserved_part(before - 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); + auto const before = balance_.fetch_add(amount, std::memory_order_acq_rel); + adaptor_->total_reserved_.add(reserved_part(before + amount) - reserved_part(before), + std::memory_order_acq_rel); + adaptor_->deallocate(stream, ptr, bytes, alignment); + } + + 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 balance before the draw-down. + std::int64_t draw_down_res(std::int64_t bytes) + { + // Nothing to enforce, so skip the compare-exchange entirely and let the balance + // go negative; the overdraft still shows up in the adaptor's `current_allocated()`. + if (is_soft()) { return balance_.fetch_sub(bytes, std::memory_order_acq_rel); } + + auto balance = balance_.load(std::memory_order_relaxed); + do { + if (bytes > balance) { + 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)); + return balance; + } + + Adaptor adaptor_; + std::int64_t const grant_; + std::size_t const overbooking_; + grant_enforcement const enforcement_; + 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. + * + * A reservation from `reserve_soft()` is not capped. It accounts identically but 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 Whether allocations may exceed the grant. + * + * @return True when the reservation was granted by `reserve_soft()`. + */ + [[nodiscard]] bool is_soft() 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/reservation_aware_resource_adaptor.hpp b/include/cucascade/memory/experimental/reservation_aware_resource_adaptor.hpp new file mode 100644 index 0000000..244e91a --- /dev/null +++ b/include/cucascade/memory/experimental/reservation_aware_resource_adaptor.hpp @@ -0,0 +1,499 @@ +/* + * 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 { + +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. +}; + +/** + * @brief Whether a reservation caps the allocations made through it. + * + * Orthogonal to `allow_overbooking`, which governs the size of the grant rather than + * what may be drawn against it. + */ +enum class grant_enforcement : bool { + STRICT, ///< Allocating past the grant throws `rmm::out_of_memory`. + SOFT, ///< Allocating past the grant is permitted and drives the balance negative. +}; + +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. + * 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 Reserve an amount of memory without capping allocations at the grant. + * + * Identical to `reserve()` in how the grant is sized and accounted, but allocations + * through the returned reservation are never refused for exceeding it. Going past the + * grant drives `memory_reservation::balance()` negative by the overdraft. The overdrawn + * bytes count only in `current_allocated()` and contribute nothing to + * `total_reserved()`, so `available()` treats the overdraft as consumed rather than as + * free space for as long as it lasts. + * + * Nothing throttles an overdraft. It is charged against the adaptor's limit but not + * bounded by it, so `available()` can go negative and stay there until the memory is + * released. Use this when the total is not known up front and a hard failure + * mid-pipeline is worse than temporarily exceeding the budget. + * + * @param size The number of bytes to reserve. + * @param overbooking_policy Whether overbooking is allowed. + * @return The reservation, which reports `memory_reservation::is_soft() == true`. + */ + [[nodiscard]] memory_reservation reserve_soft(std::size_t size, + allow_overbooking overbooking_policy); + + /** + * @brief Get the memory limit. + * + * @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; +}; + +} // namespace experimental +} // namespace memory +} // namespace cucascade + +#include diff --git a/src/memory/CMakeLists.txt b/src/memory/CMakeLists.txt index 03ba85e..6686eb6 100644 --- a/src/memory/CMakeLists.txt +++ b/src/memory/CMakeLists.txt @@ -18,21 +18,24 @@ 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/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..455b2e3 --- /dev/null +++ b/src/memory/experimental/memory_reservation.cpp @@ -0,0 +1,153 @@ +/* + * 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_); +} + +bool memory_reservation::is_soft() const noexcept +{ + return std::visit([](auto const& handle) { return handle->is_soft(); }, handle_); +} + +std::size_t memory_reservation::overbooking() const noexcept +{ + return std::visit([](auto const& handle) { return handle->overbooking(); }, handle_); +} + +} // namespace experimental +} // namespace memory +} // namespace cucascade diff --git a/src/memory/experimental/reservation_aware_resource_adaptor.cpp b/src/memory/experimental/reservation_aware_resource_adaptor.cpp new file mode 100644 index 0000000..60f5c16 --- /dev/null +++ b/src/memory/experimental/reservation_aware_resource_adaptor.cpp @@ -0,0 +1,125 @@ +/* + * 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 +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) +{ + auto const [granted, overbooking] = + this->get().reserve(size, overbooking_policy == allow_overbooking::YES); + + using impl_t = detail::memory_reservation_impl>; + return memory_reservation{ + memory_reservation::handle_variant{::cuda::mr::make_shared_resource( + *this, detail::safe_cast(granted), overbooking, grant_enforcement::STRICT)}}; +} + +template + requires ::cuda::mr::resource +memory_reservation reservation_aware_resource_adaptor::reserve_soft( + std::size_t size, allow_overbooking overbooking_policy) +{ + auto const [granted, overbooking] = + this->get().reserve(size, overbooking_policy == allow_overbooking::YES); + + using impl_t = detail::memory_reservation_impl>; + return memory_reservation{ + memory_reservation::handle_variant{::cuda::mr::make_shared_resource( + *this, detail::safe_cast(granted), overbooking, grant_enforcement::SOFT)}}; +} + +template class reservation_aware_resource_adaptor; +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 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..a01996f --- /dev/null +++ b/test/memory/test_experimental_reservation_aware_resource_adaptor.cpp @@ -0,0 +1,425 @@ +/* + * 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 + +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); + + 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); + 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_soft(1024, allow_overbooking::NO); + CHECK(res.is_soft()); + CHECK(res.balance() == 1024); + + { + // Overdrawing is permitted and shows up as a negative balance. + rmm::device_buffer buf{3072, stream, res.as_device()}; + CHECK(res.balance() == -2048); + + // The overdraft shows up as consumed memory, not as returned reserve. + CHECK(adaptor.current_allocated() == 3072); + CHECK(adaptor.total_reserved() == 0); + CHECK(adaptor.available() == limit - 3072); + } + + CHECK(res.balance() == 1024); + CHECK(adaptor.current_allocated() == 0); + CHECK(adaptor.total_reserved() == 1024); + stream.synchronize(); +} + +TEST_CASE("Overdrawn soft reservation outlived by its buffer", + "[experimental_reservation_aware][gpu]") +{ + if (!has_cuda_device()) { return; } + + rmm::cuda_stream_view stream{rmm::cuda_stream_default}; + device_adaptor adaptor{any_device_resource{rmm::mr::cuda_memory_resource{}}, limit}; + + { + rmm::device_buffer buf = [&] { + auto res = adaptor.reserve_soft(1024, allow_overbooking::NO); + return rmm::device_buffer{4096, stream, res.as_device()}; + }(); + // Only the reservation handle is gone; the buffer still holds the shared state, so + // the refund has not run yet. The grant is fully drawn, hence a zero reserve. + CHECK(adaptor.total_reserved() == 0); + CHECK(adaptor.current_allocated() == 4096); + } + + CHECK(adaptor.total_reserved() == 0); + CHECK(adaptor.current_allocated() == 0); + CHECK(adaptor.available() == limit); + stream.synchronize(); +} + +TEST_CASE("Strict reservations remain capped", "[experimental_reservation_aware][gpu]") +{ + if (!has_cuda_device()) { return; } + + rmm::cuda_stream_view stream{rmm::cuda_stream_default}; + device_adaptor adaptor{any_device_resource{rmm::mr::cuda_memory_resource{}}, limit}; + + auto res = adaptor.reserve(1024, allow_overbooking::NO); + CHECK_FALSE(res.is_soft()); + REQUIRE_THROWS_AS((rmm::device_buffer{3072, stream, res.as_device()}), rmm::out_of_memory); + CHECK(res.balance() == 1024); + stream.synchronize(); +} + +TEST_CASE("Zero-sized reservation throws on first byte", "[experimental_reservation_aware][gpu]") +{ + if (!has_cuda_device()) { return; } + + 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); + 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); + 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); + { + 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); +}