From 857ac91dce88590ecf4504b4a2c41f180f915ccb Mon Sep 17 00:00:00 2001 From: Kevin Kristensen Date: Fri, 14 Aug 2026 22:46:51 +0000 Subject: [PATCH 1/3] perf(memory): cache large pinned allocations in small_pinned_host_memory_resource Requests above the slab range previously paid a raw cudaHostAlloc / cudaFreeHost on every call; cuDF's parquet reader makes these constantly and the frees stall query execution. Serve them instead from per-bucket free lists with the same event-ordered reuse discipline as the slab path, capped at 256 MiB (constructor-tunable) with oldest-first eviction and purge-and-retry on allocation failure. Requests whose bucket exceeds the cap allocate at exact size and bypass the cache. Event pools are keyed per device so records always land on the owning context; when no event can be recorded, the large path frees instead of caching and the slab path synchronizes the freeing stream before recycling, closing a pre-existing unordered-reuse window. --- .../small_pinned_host_memory_resource.hpp | 158 +++++++-- .../small_pinned_host_memory_resource.cpp | 321 +++++++++++++++--- ...test_small_pinned_host_memory_resource.cpp | 230 ++++++++++++- 3 files changed, 643 insertions(+), 66 deletions(-) diff --git a/include/cucascade/memory/small_pinned_host_memory_resource.hpp b/include/cucascade/memory/small_pinned_host_memory_resource.hpp index cd2748f..c4497c7 100644 --- a/include/cucascade/memory/small_pinned_host_memory_resource.hpp +++ b/include/cucascade/memory/small_pinned_host_memory_resource.hpp @@ -25,6 +25,9 @@ #include #include +#include +#include +#include #include #include @@ -32,12 +35,20 @@ namespace cucascade { namespace memory { /** - * @brief A small slab allocator backed by fixed_size_host_memory_resource. + * @brief A pinned host memory allocator combining small slab pools with a bucketed reuse cache for + * large allocations. * - * Manages three pools of pinned host memory: 512 B, 1 KB, and 2 KB. - * Each pool is populated on demand by acquiring one upstream block from the - * provided fixed_size_host_memory_resource and carving it into slabs of the - * appropriate size. + * Requests up to MAX_SLAB_SIZE are served from slab pools (SLAB_SIZES) of pinned host memory. Each + * pool is populated on demand by acquiring one upstream block from the provided + * fixed_size_host_memory_resource and carving it into slabs of the appropriate size. + * + * Requests above MAX_SLAB_SIZE are rounded up to a power-of-two bucket (MIN_LARGE_BUCKET at + * minimum) and served from per-bucket free lists of previously released `cudaHostAlloc` buffers; + * misses allocate a fresh buffer. Released buffers are cached rather than freed, up to a + * configurable byte limit, sparing callers such as cuDF's parquet reader the synchronous + * `cudaHostAlloc` / `cudaFreeHost` cost on every request. Reuse in both the slab and large paths is + * ordered by CUDA events recorded on the freeing stream, so a buffer is never handed out while + * another stream may still be using it. * * Satisfies the ::cuda::mr::device_accessible and ::cuda::mr::host_accessible * properties, making it compatible with rmm::host_device_async_resource_ref @@ -56,17 +67,26 @@ namespace memory { */ class small_pinned_host_memory_resource { public: - /// Maximum allocation size handled by the slab pools. - /// Requests larger than this use pageable memory. + /// Maximum allocation size handled by the slab pools. Larger requests are served from the + /// large-allocation cache. static constexpr std::size_t MAX_SLAB_SIZE = 8192; + /// Smallest power-of-two bucket used for allocations above MAX_SLAB_SIZE. + static constexpr std::size_t MIN_LARGE_BUCKET = 16384; + + /// Default cap on the total bytes retained by the large-allocation cache. + static constexpr std::size_t DEFAULT_LARGE_CACHE_LIMIT = 256ull << 20; + /** * @brief Construct with the upstream fixed-size host memory resource. * - * @param upstream Block allocator backed by pinned host memory. Must outlive - * this object. + * @param upstream Block allocator backed by pinned host memory. Must outlive this object. + * @param large_cache_limit_bytes Maximum total bytes, measured in bucket sizes, retained by the + * large-allocation cache. */ - explicit small_pinned_host_memory_resource(fixed_size_host_memory_resource& upstream); + explicit small_pinned_host_memory_resource( + fixed_size_host_memory_resource& upstream, + std::size_t large_cache_limit_bytes = DEFAULT_LARGE_CACHE_LIMIT); small_pinned_host_memory_resource(const small_pinned_host_memory_resource&) = delete; small_pinned_host_memory_resource& operator=(const small_pinned_host_memory_resource&) = delete; @@ -78,21 +98,33 @@ class small_pinned_host_memory_resource { /** * @brief Allocate pinned memory. * - * For @p bytes <= MAX_SLAB_SIZE: rounds up to the next slab boundary - * (512 / 1 KB / 2 KB / 4 KB / 8 KB) and returns a pointer from the matching - * free list, expanding the pool from upstream if the list is empty. + * For @p bytes <= MAX_SLAB_SIZE: rounds up to the next slab boundary (512 / 1 KB / 2 KB / 4 KB / + * 8 KB) and returns a pointer from the matching free list, expanding the pool from upstream if + * the list is empty. * - * For @p bytes > MAX_SLAB_SIZE: falls back to cudaMallocHost (pinned). + * For @p bytes > MAX_SLAB_SIZE: rounds up to the power-of-two bucket and, when the bucket fits + * within the cache limit, returns a cached buffer when the bucket's free list has one, making @p + * stream wait on the buffer's ready event first. On a miss (or a never-cacheable size, which + * skips the cache lookup), allocates a fresh buffer of large_allocation_size(bytes) with + * `cudaHostAlloc(Portable | Mapped)`; if that fails, purges the entire large cache and retries + * once before throwing std::bad_alloc. */ void* allocate(::cuda::stream_ref stream, std::size_t bytes, std::size_t alignment = alignof(std::max_align_t)); /** - * @brief Return memory to the appropriate pool. + * @brief Return memory to the appropriate free list. + * + * Slabs (@p bytes <= MAX_SLAB_SIZE) are returned to the slab free list. Larger buffers are cached + * in their bucket's free list, evicting the oldest cached entries (in insertion order across + * buckets) when the total would exceed the cache limit; a buffer whose bucket alone exceeds the + * limit is freed with `cudaFreeHost` instead. Both paths record an event on @p stream so reuse + * waits for pending work on the buffer. When no event can be recorded, @p stream is synchronized + * instead (best effort); a large buffer is then freed rather than cached and a slab is cached + * carrying no pending work, so a cached entry never carries pending work that reuse cannot order + * against. * - * Slabs (@p bytes <= MAX_SLAB_SIZE) are returned to the free list. - * Pinned allocations (@p bytes > MAX_SLAB_SIZE) are freed via cudaFreeHost. * @p bytes must equal the value passed to the corresponding allocate. */ void deallocate(::cuda::stream_ref stream, @@ -115,6 +147,9 @@ class small_pinned_host_memory_resource { rmm::cuda_stream_default.synchronize_no_throw(); } + /// Total bytes currently held in the large-allocation cache, measured in bucket sizes. + [[nodiscard]] std::size_t large_cache_bytes() const; + bool operator==(small_pinned_host_memory_resource const& other) const noexcept; /** @@ -141,32 +176,105 @@ class small_pinned_host_memory_resource { /// Returns the index into SLAB_SIZES of the smallest slab >= bytes. static std::size_t slab_index_for(std::size_t bytes) noexcept; + /// Returns the bucket for a request above MAX_SLAB_SIZE: the smallest power of two >= bytes, no + /// smaller than MIN_LARGE_BUCKET. The bucket is used for cache keying and eviction accounting; + /// the physical size is large_allocation_size(bytes), equal to the bucket only for cacheable + /// requests. allocate and deallocate both derive the bucket from the request size, so the pairing + /// is deterministic. + static std::size_t large_bucket_size_for(std::size_t bytes) noexcept; + + /// Returns the physical size cudaHostAlloc is asked for on a cache miss: the bucket when it fits + /// within the cache limit, otherwise exactly @p bytes. A never-cacheable allocation gains nothing + /// from bucket rounding and must not overshoot pinned memory (up to 2x for sizes just past a + /// bucket boundary). Reads only the immutable cache limit, so no lock is needed. + [[nodiscard]] std::size_t large_allocation_size(std::size_t bytes) const noexcept; + /// Populate the free list for slab @p idx by acquiring one upstream block. /// Must be called with mutex_ held. void expand_pool_locked(std::size_t slab_idx); + /// A CUDA event paired with the device whose context owns it. cudaEventRecord requires the event + /// and the stream to share a CUDA context, so pooled events are segregated by device and an event + /// is only ever recorded on a stream of its own device. @c device is meaningful only when @c + /// handle is non-null. + struct device_event { + cudaEvent_t handle{nullptr}; + int device{-1}; + }; + /// A free slab plus, when it was just deallocated, a CUDA event recorded on /// the freeing stream. Reusing the slab must wait on this event so an /// in-flight async H2D copy that still reads the slab (e.g. cuDF's parquet /// stats min/max buffers) completes before another stream overwrites it. - /// @c ready_event is null for freshly-carved slabs that were never used. + /// A null @c ready.handle means the slab carries no pending work: it was + /// freshly carved and never used, or the freeing stream was synchronized + /// before the slab was cached. struct free_slab { void* ptr; - cudaEvent_t ready_event; + device_event ready; }; - /// Borrow a timing-disabled CUDA event (recycled from @c event_pool_ or newly - /// created). Returns null if event creation fails. Must hold @c mutex_. - cudaEvent_t acquire_event_locked(); + /// A released large buffer held for reuse. Like @c free_slab, @c ready captures the freeing + /// stream's last use of the buffer; unlike a slab, a large buffer is never cached without a + /// recorded event (deallocate frees it instead), so @c ready.handle is non-null for every cached + /// entry. @c sequence orders entries across buckets so eviction can drop the oldest first. + struct large_cache_entry { + void* ptr; + device_event ready; + std::uint64_t sequence; + }; + + /// Borrow a timing-disabled CUDA event owned by the calling thread's current device (recycled + /// from @c event_pools_ or newly created). Returns a null handle when the device query or event + /// creation fails; the caller then falls back to freeing (large path) or synchronizing (slab + /// path) instead of caching with ordering. Must hold @c mutex_. + device_event acquire_event_locked(); + + /// Return an event to its device's pool in @c event_pools_ for reuse. Must hold @c mutex_. + void release_event_locked(device_event event) noexcept; + + /// Pop the oldest cached buffer of @p bucket and make @p stream wait on its ready event, + /// recycling the event. Returns null when the bucket has no cached entries. Must hold @c mutex_. + void* try_take_cached_large_locked(std::size_t bucket, ::cuda::stream_ref stream); + + /// Remove the oldest cached large buffer across all buckets from the bookkeeping and return it + /// with its ready event still attached, for the caller to pass to sync_and_free_large_victim + /// outside the lock and then recycle the event. Returns a null-ptr entry when the cache is empty. + /// Must hold @c mutex_. + large_cache_entry evict_oldest_large_locked() noexcept; - /// Return an event to @c event_pool_ for reuse. Must hold @c mutex_. - void release_event_locked(cudaEvent_t event) noexcept; + /// Remove every cached large buffer from the bookkeeping and return them with their ready events + /// still attached, for the caller to pass to sync_and_free_large_victim outside the lock and then + /// recycle the events. Must hold @c mutex_. + std::vector purge_large_cache_locked(); + + /// Wait for a victim's recorded work and free its buffer. The ready event may have been recorded + /// on any device's stream, so the buffer is not unpinned until the event has completed + /// (best-effort: a failed or null event skips the wait). Must be called without @c mutex_ held; + /// the caller recycles the event afterwards. + static void sync_and_free_large_victim(large_cache_entry const& victim) noexcept; fixed_size_host_memory_resource& upstream_; mutable std::mutex mutex_; std::array, 5> free_lists_{}; - std::vector event_pool_; + + /// Idle recycled events keyed by the device that owns them. Events are created in the calling + /// thread's current context, so acquire_event_locked keys by cudaGetDevice() at acquire time; + /// this avoids cudaSetDevice churn, and in deployment the calling thread's current device is its + /// stream's device (the NUMA dispatcher routes by cudaGetDevice()). Should a caller ever pass a + /// stream of another device, the record fails and the safe fallback engages, so correctness never + /// depends on the key, only cache-hit rate. + std::map> event_pools_; + std::vector owned_allocations_; + + /// Large-allocation cache: per-bucket free lists keyed by bucket size, the current total in + /// bucket bytes, the retention cap, and a monotonic counter stamping insertion order for + /// eviction. + std::map> large_cache_; + std::size_t large_cache_bytes_ = 0; + std::size_t const large_cache_limit_bytes_; + std::uint64_t large_cache_sequence_ = 0; }; static_assert(::cuda::mr::resource_with +#include +#include #include +#include +#include #include +#include namespace cucascade { namespace memory { small_pinned_host_memory_resource::small_pinned_host_memory_resource( - fixed_size_host_memory_resource& upstream) - : upstream_(upstream) + fixed_size_host_memory_resource& upstream, std::size_t large_cache_limit_bytes) + : upstream_(upstream), large_cache_limit_bytes_(large_cache_limit_bytes) { } @@ -35,44 +40,73 @@ small_pinned_host_memory_resource::~small_pinned_host_memory_resource() { // owned_allocations_ destructor returns upstream blocks to the free list. // free_lists_ entries are raw pointers into those blocks; no individual cleanup needed. + // Cached large buffers come straight from cudaHostAlloc, so each is freed here. // Destroy every CUDA event we own — both idle (pooled) and still attached to a - // free slab that was never re-allocated. - for (auto& event : event_pool_) { - if (event != nullptr) { CUCASCADE_ASSERT_CUDA_SUCCESS(::cudaEventDestroy(event)); } + // free slab or cached large buffer that was never re-allocated. The events may + // be owned by multiple devices' contexts; cudaEventDestroy carries no + // same-device precondition, so destruction order and the current device are + // irrelevant. + for (auto& pool : event_pools_) { + for (auto& event : pool.second) { + if (event != nullptr) { CUCASCADE_ASSERT_CUDA_SUCCESS(::cudaEventDestroy(event)); } + } } for (auto& list : free_lists_) { for (auto& slab : list) { - if (slab.ready_event != nullptr) { - CUCASCADE_ASSERT_CUDA_SUCCESS(::cudaEventDestroy(slab.ready_event)); + if (slab.ready.handle != nullptr) { + CUCASCADE_ASSERT_CUDA_SUCCESS(::cudaEventDestroy(slab.ready.handle)); + } + } + } + // Destruction assumes no further allocate/deallocate calls, but work already recorded on a + // cached entry's ready event may still be draining, so each buffer goes through the same + // wait-then-free as eviction and purge before its event is destroyed. + for (auto& bucket : large_cache_) { + for (auto& entry : bucket.second) { + sync_and_free_large_victim(entry); + if (entry.ready.handle != nullptr) { + CUCASCADE_ASSERT_CUDA_SUCCESS(::cudaEventDestroy(entry.ready.handle)); } } } } -cudaEvent_t small_pinned_host_memory_resource::acquire_event_locked() +small_pinned_host_memory_resource::device_event +small_pinned_host_memory_resource::acquire_event_locked() { - if (!event_pool_.empty()) { - cudaEvent_t event = event_pool_.back(); - event_pool_.pop_back(); - return event; + int device = -1; + if (::cudaGetDevice(&device) != cudaSuccess) { + // Best effort: a null handle routes the caller to its no-event fallback, and allocation must + // never throw here. Clear the sticky error and carry on. + (void)::cudaGetLastError(); + return {}; + } + auto& pool = event_pools_[device]; + if (!pool.empty()) { + cudaEvent_t event = pool.back(); + pool.pop_back(); + return {event, device}; } cudaEvent_t event = nullptr; // Timing is not needed; disabling it makes record/wait cheaper. if (::cudaEventCreateWithFlags(&event, cudaEventDisableTiming) != cudaSuccess) { - // Best effort: without an event this deallocation loses stream ordering, but - // allocation must never throw here. Clear the sticky error and carry on. + // Best effort, as above: the caller handles the null handle; allocation must never throw + // here. Clear the sticky error and carry on. (void)::cudaGetLastError(); - return nullptr; + return {}; } - return event; + return {event, device}; } -void small_pinned_host_memory_resource::release_event_locked(cudaEvent_t event) noexcept +void small_pinned_host_memory_resource::release_event_locked(device_event event) noexcept { - if (event != nullptr) { event_pool_.push_back(event); } + if (event.handle != nullptr) { + assert(event.device >= 0); + event_pools_[event.device].push_back(event.handle); + } } -void* small_pinned_host_memory_resource::allocate([[maybe_unused]] cuda::stream_ref stream, +void* small_pinned_host_memory_resource::allocate(cuda::stream_ref stream, std::size_t bytes, [[maybe_unused]] std::size_t alignment) { @@ -87,10 +121,48 @@ void* small_pinned_host_memory_resource::allocate([[maybe_unused]] cuda::stream_ // directly from GPU kernels (e.g. detect_malformed_pages), so returning pageable memory // here would cause cudaErrorIllegalAddress. if (bytes > MAX_SLAB_SIZE) { + std::size_t const bucket = large_bucket_size_for(bytes); + bool const cacheable = bucket <= large_cache_limit_bytes_; + std::size_t const alloc_bytes = large_allocation_size(bytes); + if (cacheable) { + std::lock_guard lock(mutex_); + if (void* cached = try_take_cached_large_locked(bucket, stream)) { return cached; } + } void* ptr = nullptr; // Portable + Mapped — see numa_region_pinned_host_allocator.cpp comment. - auto err = ::cudaHostAlloc(&ptr, bytes, cudaHostAllocPortable | cudaHostAllocMapped); - if (err != cudaSuccess) { throw std::bad_alloc{}; } + auto err = ::cudaHostAlloc(&ptr, alloc_bytes, cudaHostAllocPortable | cudaHostAllocMapped); + if (err == cudaSuccess) { return ptr; } + // Clear the sticky error so a successful retry does not leave cudaGetLastError consumers + // seeing a stale allocation failure. + (void)::cudaGetLastError(); + // Cached buffers are the only pinned memory this class can give back under pressure: + // release them all and retry once (even a never-cacheable request benefits, since purging + // frees the pinned memory its retry needs). Victims are synchronized and freed outside the + // lock so slab traffic does not stall behind the waits; their events are recycled in one + // batch under the re-taken lock. + std::vector purged; + { + std::lock_guard lock(mutex_); + purged = purge_large_cache_locked(); + } + for (auto const& victim : purged) { + sync_and_free_large_victim(victim); + } + { + std::lock_guard lock(mutex_); + for (auto const& victim : purged) { + release_event_locked(victim.ready); + } + // A racing deallocate may have cached a matching buffer while we freed; take it if so. + if (cacheable) { + if (void* cached = try_take_cached_large_locked(bucket, stream)) { return cached; } + } + } + err = ::cudaHostAlloc(&ptr, alloc_bytes, cudaHostAllocPortable | cudaHostAllocMapped); + if (err != cudaSuccess) { + (void)::cudaGetLastError(); + throw std::bad_alloc{}; + } return ptr; } @@ -99,41 +171,115 @@ void* small_pinned_host_memory_resource::allocate([[maybe_unused]] cuda::stream_ if (free_lists_[idx].empty()) { expand_pool_locked(idx); } free_slab slab = free_lists_[idx].back(); free_lists_[idx].pop_back(); - // If this slab was recently deallocated, its ready_event captures the freeing + // If this slab was recently deallocated, its ready event captures the freeing // stream's last use (e.g. an in-flight async H2D copy still reading the slab). // Make the reusing stream wait for it so we cannot overwrite the slab before // that copy completes. Recording the event again (on a later deallocate) does - // not disturb this already-enqueued wait, so the event is safe to recycle. - if (slab.ready_event != nullptr) { - CUCASCADE_ASSERT_CUDA_SUCCESS(::cudaStreamWaitEvent(stream.get(), slab.ready_event, 0)); - release_event_locked(slab.ready_event); + // not disturb this already-enqueued wait, so the event is safe to recycle. A + // null handle means the slab carries no pending work, so no wait is needed. + if (slab.ready.handle != nullptr) { + CUCASCADE_ASSERT_CUDA_SUCCESS(::cudaStreamWaitEvent(stream.get(), slab.ready.handle, 0)); + release_event_locked(slab.ready); } return slab.ptr; } -void small_pinned_host_memory_resource::deallocate([[maybe_unused]] cuda::stream_ref stream, +void small_pinned_host_memory_resource::deallocate(cuda::stream_ref stream, void* ptr, std::size_t bytes, [[maybe_unused]] std::size_t alignment) noexcept { if (ptr == nullptr || bytes == 0) { return; } if (bytes > MAX_SLAB_SIZE) { - ::cudaFreeHost(ptr); - return; + std::size_t const bucket = large_bucket_size_for(bytes); + if (bucket > large_cache_limit_bytes_) { + // Too big to ever cache: free directly. Any pending work on the buffer was issued by the + // caller on its own device, which is the case cudaFreeHost's implicit synchronization + // covers; no ready event has been recorded for this buffer. + CUCASCADE_ASSERT_CUDA_SUCCESS(::cudaFreeHost(ptr)); + return; + } + // Evict until the incoming bucket fits, then cache it. A victim's pending work may live on + // any device's stream, so each iteration unhooks one victim under the lock, waits on its + // ready event and frees it outside the lock (slab traffic must not stall behind the wait), + // and then recycles the event. + while (true) { + large_cache_entry victim{nullptr, {}, 0}; + bool free_instead_of_caching = false; + { + std::lock_guard lock(mutex_); + if (large_cache_bytes_ + bucket <= large_cache_limit_bytes_) { + // As in the slab path below, the recorded event defers reuse until pending work on + // the freeing stream completes. + device_event event = acquire_event_locked(); + if (event.handle != nullptr && + ::cudaEventRecord(event.handle, stream.get()) != cudaSuccess) { + (void)::cudaGetLastError(); + release_event_locked(event); + event = {}; + } + if (event.handle != nullptr) { + large_cache_[bucket].push_back(large_cache_entry{ptr, event, large_cache_sequence_++}); + large_cache_bytes_ += bucket; + return; + } + // Reuse ordering is carried solely by the recorded event, so a buffer whose event + // could not be recorded is freed (outside the lock, below) instead of cached. This + // branch is rare because pooled events are segregated by device: it fires only on + // event-creation failure or a stream that does not belong to the caller's current + // device. + free_instead_of_caching = true; + } else { + victim = evict_oldest_large_locked(); + } + } + if (free_instead_of_caching) { + // With no event recorded on the freeing stream, ordering comes from draining the stream + // before the free, exactly as in the slab fallback below. The stream may belong to + // another device's context, which cudaFreeHost's implicit synchronization is not + // documented to cover. Best effort on failure: clear the sticky error and free anyway. + if (::cudaStreamSynchronize(stream.get()) != cudaSuccess) { (void)::cudaGetLastError(); } + CUCASCADE_ASSERT_CUDA_SUCCESS(::cudaFreeHost(ptr)); + return; + } + if (victim.ptr == nullptr) { + // Unreachable while the bookkeeping is consistent (bucket <= limit means an empty + // cache fits); free the buffer rather than loop forever. + CUCASCADE_ASSERT_CUDA_SUCCESS(::cudaFreeHost(ptr)); + return; + } + sync_and_free_large_victim(victim); + if (victim.ready.handle != nullptr) { + std::lock_guard lock(mutex_); + release_event_locked(victim.ready); + } + } } std::size_t idx = slab_index_for(bytes); - std::lock_guard lock(mutex_); - // Record an event on the freeing stream so a future reuse of this slab can wait - // for any still-pending work on it (the async H2D copy in cuDF's stats filter). - // Best effort: a null event just means this slab is recycled without ordering. - cudaEvent_t event = acquire_event_locked(); - if (event != nullptr && ::cudaEventRecord(event, stream.get()) != cudaSuccess) { - (void)::cudaGetLastError(); - release_event_locked(event); - event = nullptr; + { + std::lock_guard lock(mutex_); + // Record an event on the freeing stream so a future reuse of this slab can wait + // for any still-pending work on it (the async H2D copy in cuDF's stats filter). + device_event event = acquire_event_locked(); + if (event.handle != nullptr) { + if (::cudaEventRecord(event.handle, stream.get()) == cudaSuccess) { + free_lists_[idx].push_back(free_slab{ptr, event}); + return; + } + (void)::cudaGetLastError(); + release_event_locked(event); + } } - free_lists_[idx].push_back(free_slab{ptr, event}); + // No event could be recorded, so drain the freeing stream before caching: a null ready handle + // promises the slab carries no pending work, and synchronizing here is what keeps that promise. + // Best effort on failure: a stream broken enough to fail synchronize is a context where ordering + // is already lost, and leaking the slab would be the only alternative. + if (::cudaStreamSynchronize(stream.get()) != cudaSuccess) { (void)::cudaGetLastError(); } + // The push must happen after the synchronize completes; pushing first would let a racing + // allocate hand the slab out before the freeing stream drains. + std::lock_guard lock(mutex_); + free_lists_[idx].push_back(free_slab{ptr, {}}); } bool small_pinned_host_memory_resource::operator==( @@ -150,6 +296,101 @@ std::size_t small_pinned_host_memory_resource::slab_index_for(std::size_t bytes) return SLAB_SIZES.size() - 1; } +std::size_t small_pinned_host_memory_resource::large_bucket_size_for(std::size_t bytes) noexcept +{ + if (bytes <= MIN_LARGE_BUCKET) { return MIN_LARGE_BUCKET; } + // bit_ceil is undefined when the next power of two is unrepresentable. Such a request cannot + // be satisfied anyway, so pass it through unrounded and let cudaHostAlloc reject it. + constexpr std::size_t max_bucket = std::size_t{1} + << (std::numeric_limits::digits - 1); + if (bytes > max_bucket) { return bytes; } + return std::bit_ceil(bytes); +} + +std::size_t small_pinned_host_memory_resource::large_allocation_size( + std::size_t bytes) const noexcept +{ + std::size_t const bucket = large_bucket_size_for(bytes); + return bucket <= large_cache_limit_bytes_ ? bucket : bytes; +} + +std::size_t small_pinned_host_memory_resource::large_cache_bytes() const +{ + std::lock_guard lock(mutex_); + return large_cache_bytes_; +} + +void* small_pinned_host_memory_resource::try_take_cached_large_locked(std::size_t bucket, + ::cuda::stream_ref stream) +{ + auto it = large_cache_.find(bucket); + if (it == large_cache_.end()) { return nullptr; } + large_cache_entry entry = it->second.front(); + it->second.pop_front(); + if (it->second.empty()) { large_cache_.erase(it); } + large_cache_bytes_ -= bucket; + // Same reuse discipline as the slab path: make the reusing stream wait for the freeing + // stream's last use of this buffer before it can be overwritten. + if (entry.ready.handle != nullptr) { + CUCASCADE_ASSERT_CUDA_SUCCESS(::cudaStreamWaitEvent(stream.get(), entry.ready.handle, 0)); + release_event_locked(entry.ready); + } + return entry.ptr; +} + +small_pinned_host_memory_resource::large_cache_entry +small_pinned_host_memory_resource::evict_oldest_large_locked() noexcept +{ + // Buckets never hold an empty deque (removals erase drained buckets), so every front() is + // that bucket's oldest entry and the global oldest is the minimum sequence across fronts. + auto oldest = large_cache_.end(); + for (auto it = large_cache_.begin(); it != large_cache_.end(); ++it) { + if (oldest == large_cache_.end() || + it->second.front().sequence < oldest->second.front().sequence) { + oldest = it; + } + } + if (oldest == large_cache_.end()) { return large_cache_entry{nullptr, {}, 0}; } + std::size_t const bucket = oldest->first; + large_cache_entry entry = oldest->second.front(); + oldest->second.pop_front(); + if (oldest->second.empty()) { large_cache_.erase(oldest); } + large_cache_bytes_ -= bucket; + return entry; +} + +std::vector +small_pinned_host_memory_resource::purge_large_cache_locked() +{ + std::vector entries; + std::size_t count = 0; + for (auto& bucket : large_cache_) { + count += bucket.second.size(); + } + entries.reserve(count); + for (auto& bucket : large_cache_) { + for (auto& entry : bucket.second) { + entries.push_back(entry); + } + } + large_cache_.clear(); + large_cache_bytes_ = 0; + return entries; +} + +void small_pinned_host_memory_resource::sync_and_free_large_victim( + large_cache_entry const& victim) noexcept +{ + // The ready event may target any device's stream; wait for it before unpinning the pages so + // an in-flight DMA cannot read freed memory. Best effort: on a failed wait, clear the sticky + // error and free anyway. + if (victim.ready.handle != nullptr && + ::cudaEventSynchronize(victim.ready.handle) != cudaSuccess) { + (void)::cudaGetLastError(); + } + CUCASCADE_ASSERT_CUDA_SUCCESS(::cudaFreeHost(victim.ptr)); +} + void small_pinned_host_memory_resource::expand_pool_locked(std::size_t slab_idx) { // Acquire one upstream block and carve it into slabs. @@ -161,7 +402,7 @@ void small_pinned_host_memory_resource::expand_pool_locked(std::size_t slab_idx) for (std::byte* block : allocation->get_blocks()) { for (std::size_t i = 0; i < num_slabs; ++i) { // Freshly-carved slabs were never used, so they carry no pending-work event. - free_lists_[slab_idx].push_back(free_slab{block + i * slab_size, nullptr}); + free_lists_[slab_idx].push_back(free_slab{block + i * slab_size, {}}); } } owned_allocations_.push_back(std::move(allocation)); diff --git a/test/memory/test_small_pinned_host_memory_resource.cpp b/test/memory/test_small_pinned_host_memory_resource.cpp index 349adba..9ff8143 100644 --- a/test/memory/test_small_pinned_host_memory_resource.cpp +++ b/test/memory/test_small_pinned_host_memory_resource.cpp @@ -17,6 +17,7 @@ #include +#include #include #include @@ -24,6 +25,7 @@ #include #include +#include #include #include #include @@ -95,7 +97,7 @@ TEST_CASE("Sub-slab sizes round up correctly", "[small_pinned]") f.slab_mr.deallocate(rmm::cuda_stream_view{}, p3, 4097); } -TEST_CASE("Large allocation falls back to malloc", "[small_pinned]") +TEST_CASE("Large allocation is served by the bucketed pinned path", "[small_pinned]") { test_fixture f; constexpr std::size_t big = small_pinned_host_memory_resource::MAX_SLAB_SIZE + 1; @@ -272,3 +274,229 @@ TEST_CASE("Large allocations do not interfere with slab pool", "[small_pinned]") f.slab_mr.deallocate(rmm::cuda_stream_view{}, small1, small_size); f.slab_mr.deallocate(rmm::cuda_stream_view{}, small2, small_size); } + +TEST_CASE("large_bucket_size_for rounds to power-of-two buckets", "[small_pinned]") +{ + using mr = small_pinned_host_memory_resource; + REQUIRE(mr::large_bucket_size_for(mr::MAX_SLAB_SIZE + 1) == mr::MIN_LARGE_BUCKET); + REQUIRE(mr::large_bucket_size_for(16384) == 16384); + REQUIRE(mr::large_bucket_size_for(16385) == 32768); + REQUIRE(mr::large_bucket_size_for(100000) == 131072); + REQUIRE(mr::large_bucket_size_for(std::size_t{1} << 20) == std::size_t{1} << 20); +} + +TEST_CASE("large_allocation_size rounds only cacheable requests", "[small_pinned]") +{ + // allocate sizes its cudaHostAlloc calls with this same helper, so these checks pin down the + // physical size behavior: bucket rounding for cacheable requests, the exact request otherwise. + test_fixture f; + using mr_t = small_pinned_host_memory_resource; + REQUIRE(f.slab_mr.large_allocation_size(100000) == 131072); + REQUIRE(f.slab_mr.large_allocation_size(mr_t::MIN_LARGE_BUCKET) == mr_t::MIN_LARGE_BUCKET); + REQUIRE(f.slab_mr.large_allocation_size(mr_t::DEFAULT_LARGE_CACHE_LIMIT) == + mr_t::DEFAULT_LARGE_CACHE_LIMIT); + REQUIRE(f.slab_mr.large_allocation_size(mr_t::DEFAULT_LARGE_CACHE_LIMIT + 1) == + mr_t::DEFAULT_LARGE_CACHE_LIMIT + 1); + + small_pinned_host_memory_resource tiny{f.upstream, 64 * 1024}; + REQUIRE(tiny.large_allocation_size(9000) == 16384); + REQUIRE(tiny.large_allocation_size(100000) == 100000); +} + +TEST_CASE("Large allocations are cached and reused", "[small_pinned]") +{ + test_fixture f; + constexpr std::size_t big = 100000; // rounds up to a 128 KB bucket + + auto* slab = f.slab_mr.allocate(rmm::cuda_stream_view{}, 512); + auto* p1 = f.slab_mr.allocate(rmm::cuda_stream_view{}, big); + REQUIRE(p1 != nullptr); + REQUIRE(p1 != slab); + std::memset(p1, 0xEF, big); + f.slab_mr.deallocate(rmm::cuda_stream_view{}, p1, big); + + // The freed buffer is cached, so an allocation of the same size gets it back. + auto* p2 = f.slab_mr.allocate(rmm::cuda_stream_view{}, big); + REQUIRE(p2 == p1); + f.slab_mr.deallocate(rmm::cuda_stream_view{}, p2, big); + f.slab_mr.deallocate(rmm::cuda_stream_view{}, slab, 512); +} + +TEST_CASE("Large cache serves within a bucket but not across buckets", "[small_pinned]") +{ + test_fixture f; + + // 12 KB and 40 KB round to different buckets (16 KB and 64 KB): no cross-serving. The 12 KB + // buffer stays cached (still allocated), so the 64 KB miss cannot alias it. + auto* p16 = f.slab_mr.allocate(rmm::cuda_stream_view{}, 12 * 1024); + f.slab_mr.deallocate(rmm::cuda_stream_view{}, p16, 12 * 1024); + auto* p64 = f.slab_mr.allocate(rmm::cuda_stream_view{}, 40 * 1024); + REQUIRE(p64 != p16); + f.slab_mr.deallocate(rmm::cuda_stream_view{}, p64, 40 * 1024); + + // 33 KB and 40 KB round to the same 64 KB bucket: the cached buffer is reused. + auto* p1 = f.slab_mr.allocate(rmm::cuda_stream_view{}, 33 * 1024); + f.slab_mr.deallocate(rmm::cuda_stream_view{}, p1, 33 * 1024); + auto* p2 = f.slab_mr.allocate(rmm::cuda_stream_view{}, 40 * 1024); + REQUIRE(p2 == p1); + f.slab_mr.deallocate(rmm::cuda_stream_view{}, p2, 40 * 1024); +} + +TEST_CASE("large_cache_bytes reflects bucket-size accounting", "[small_pinned]") +{ + test_fixture f; + REQUIRE(f.slab_mr.large_cache_bytes() == 0); + + // 100000 bytes occupy a 128 KB bucket; the cache counts the bucket, not the request. + constexpr std::size_t bytes = 100000; + auto* p = f.slab_mr.allocate(rmm::cuda_stream_view{}, bytes); + REQUIRE(f.slab_mr.large_cache_bytes() == 0); // live buffers are not cached + f.slab_mr.deallocate(rmm::cuda_stream_view{}, p, bytes); + REQUIRE(f.slab_mr.large_cache_bytes() == 128 * 1024); + + // A size one past a bucket boundary lands in the next bucket. + auto* q = f.slab_mr.allocate(rmm::cuda_stream_view{}, 16 * 1024 + 1); + f.slab_mr.deallocate(rmm::cuda_stream_view{}, q, 16 * 1024 + 1); + REQUIRE(f.slab_mr.large_cache_bytes() == 128 * 1024 + 32 * 1024); + + // Cache hits remove the bucket from the total. + auto* r = f.slab_mr.allocate(rmm::cuda_stream_view{}, bytes); + REQUIRE(r == p); + REQUIRE(f.slab_mr.large_cache_bytes() == 32 * 1024); + f.slab_mr.deallocate(rmm::cuda_stream_view{}, r, bytes); +} + +TEST_CASE("Large cache respects its cap and evicts oldest entries", "[small_pinned]") +{ + test_fixture f; + constexpr std::size_t cap = 64 * 1024; + small_pinned_host_memory_resource mr{f.upstream, cap}; + + constexpr std::size_t bytes = 9000; // bucket = 16 KB, so the cap holds four buffers + std::array ptrs{}; + for (auto& p : ptrs) { + p = mr.allocate(rmm::cuda_stream_view{}, bytes); + REQUIRE(p != nullptr); + } + for (auto* p : ptrs) { + mr.deallocate(rmm::cuda_stream_view{}, p, bytes); + REQUIRE(mr.large_cache_bytes() <= cap); + } + // Four 16 KB buckets fill the cap; caching the fifth evicted the oldest entry (ptrs[0]). + REQUIRE(mr.large_cache_bytes() == cap); + + std::set reused; + for (int i = 0; i < 4; ++i) { + reused.insert(mr.allocate(rmm::cuda_stream_view{}, bytes)); + } + REQUIRE(mr.large_cache_bytes() == 0); + REQUIRE(reused == std::set{ptrs[1], ptrs[2], ptrs[3], ptrs[4]}); + for (auto* p : reused) { + mr.deallocate(rmm::cuda_stream_view{}, p, bytes); + } +} + +TEST_CASE("Buffers larger than the whole cap are freed, not cached", "[small_pinned]") +{ + test_fixture f; + small_pinned_host_memory_resource mr{f.upstream, 64 * 1024}; + + // Seed the cache so we can also verify the oversized free evicts nothing. + auto* seeded = mr.allocate(rmm::cuda_stream_view{}, 9000); + mr.deallocate(rmm::cuda_stream_view{}, seeded, 9000); + auto const cached_before = mr.large_cache_bytes(); + REQUIRE(cached_before == 16 * 1024); + + constexpr std::size_t huge = 128 * 1024; // bucket exceeds the whole 64 KB cap + auto* p = mr.allocate(rmm::cuda_stream_view{}, huge); + REQUIRE(p != nullptr); + mr.deallocate(rmm::cuda_stream_view{}, p, huge); + REQUIRE(mr.large_cache_bytes() == cached_before); +} + +TEST_CASE("Never-cacheable sizes round trip without polluting the cache", "[small_pinned]") +{ + // A non-power-of-two size whose bucket exceeds the whole cap is allocated at exactly the + // requested size and freed on deallocate rather than cached. + test_fixture f; + small_pinned_host_memory_resource mr{f.upstream, 64 * 1024}; + + constexpr std::size_t bytes = 100000; // 128 KB bucket, beyond the 64 KB cap + auto* p = mr.allocate(rmm::cuda_stream_view{}, bytes); + REQUIRE(p != nullptr); + std::memset(p, 0x5A, bytes); // the full requested size must be usable + mr.deallocate(rmm::cuda_stream_view{}, p, bytes); + REQUIRE(mr.large_cache_bytes() == 0); + + auto* q = mr.allocate(rmm::cuda_stream_view{}, bytes); + REQUIRE(q != nullptr); + mr.deallocate(rmm::cuda_stream_view{}, q, bytes); +} + +TEST_CASE("Slab and large reuse work on a real stream", "[small_pinned]") +{ + // Exercises the deallocate-time event record and allocate-time wait on a genuine + // (non-default) stream, driving the device-keyed event pool end to end. + test_fixture f; + rmm::cuda_stream stream; + + constexpr std::size_t slab_bytes = 2048; + auto* s1 = f.slab_mr.allocate(stream.view(), slab_bytes); + REQUIRE(s1 != nullptr); + std::memset(s1, 0x11, slab_bytes); + f.slab_mr.deallocate(stream.view(), s1, slab_bytes); + auto* s2 = f.slab_mr.allocate(stream.view(), slab_bytes); + REQUIRE(s2 == s1); + f.slab_mr.deallocate(stream.view(), s2, slab_bytes); + + constexpr std::size_t big_bytes = 100000; + auto* p1 = f.slab_mr.allocate(stream.view(), big_bytes); + REQUIRE(p1 != nullptr); + std::memset(p1, 0x22, big_bytes); + f.slab_mr.deallocate(stream.view(), p1, big_bytes); + auto* p2 = f.slab_mr.allocate(stream.view(), big_bytes); + REQUIRE(p2 == p1); + f.slab_mr.deallocate(stream.view(), p2, big_bytes); + stream.synchronize(); +} + +TEST_CASE("Destruction with a populated large cache is safe", "[small_pinned]") +{ + test_fixture f; + { + // Deallocate -> allocate -> deallocate cycles on both paths recycle events through the idle + // pool, so the events destroyed with this resource include recycled ones, not only fresh + // ones, attached to a free slab as well as to cached large buffers. + small_pinned_host_memory_resource mr{f.upstream}; + auto* s = mr.allocate(rmm::cuda_stream_view{}, 512); + mr.deallocate(rmm::cuda_stream_view{}, s, 512); + auto* s2 = mr.allocate(rmm::cuda_stream_view{}, 512); + REQUIRE(s2 == s); + mr.deallocate(rmm::cuda_stream_view{}, s2, 512); + + auto* p1 = mr.allocate(rmm::cuda_stream_view{}, 10 * 1024); + auto* p2 = mr.allocate(rmm::cuda_stream_view{}, 100 * 1024); + mr.deallocate(rmm::cuda_stream_view{}, p1, 10 * 1024); + mr.deallocate(rmm::cuda_stream_view{}, p2, 100 * 1024); + auto* p3 = mr.allocate(rmm::cuda_stream_view{}, 10 * 1024); + REQUIRE(p3 == p1); + mr.deallocate(rmm::cuda_stream_view{}, p3, 10 * 1024); + REQUIRE(mr.large_cache_bytes() == 16 * 1024 + 128 * 1024); + } + { + // A 32 KB insertion into a cap-full cache of two 16 KB entries evicts both victims but + // re-acquires only one event, so this resource is destroyed while its pool holds an idle + // recycled event alongside the events still attached to a free slab and a cached buffer. + small_pinned_host_memory_resource mr{f.upstream, 32 * 1024}; + auto* s = mr.allocate(rmm::cuda_stream_view{}, 512); + mr.deallocate(rmm::cuda_stream_view{}, s, 512); + auto* a = mr.allocate(rmm::cuda_stream_view{}, 9000); + auto* b = mr.allocate(rmm::cuda_stream_view{}, 9000); + auto* c = mr.allocate(rmm::cuda_stream_view{}, 20 * 1024); + mr.deallocate(rmm::cuda_stream_view{}, a, 9000); + mr.deallocate(rmm::cuda_stream_view{}, b, 9000); + mr.deallocate(rmm::cuda_stream_view{}, c, 20 * 1024); + REQUIRE(mr.large_cache_bytes() == 32 * 1024); + } + SUCCEED("resources destroyed with cached buffers, attached events, and pooled idle events"); +} From de27f152e10aab6f154b6a797b3e71578be2b0dc Mon Sep 17 00:00:00 2001 From: Kevin Kristensen Date: Fri, 14 Aug 2026 23:49:20 +0000 Subject: [PATCH 2/3] fix(memory): make pinned host cache stream-safe --- CLAUDE.md | 2 +- include/cucascade/memory/common.hpp | 4 + .../memory/null_device_memory_resource.hpp | 4 + .../numa_region_pinned_host_allocator.hpp | 4 + .../small_pinned_host_memory_resource.hpp | 218 +++++++++--------- .../small_pinned_host_memory_resource.cpp | 96 +++----- ...est_reservation_aware_resource_adaptor.cpp | 4 + ...test_small_pinned_host_memory_resource.cpp | 1 - 8 files changed, 153 insertions(+), 180 deletions(-) diff --git a/CLAUDE.md b/CLAUDE.md index 0750d2e..1745f2c 100644 --- a/CLAUDE.md +++ b/CLAUDE.md @@ -148,7 +148,7 @@ Performance optimization of cuCascade's disk I/O backends (GDS and kvikIO) to ap #include #include // cuDF #include // RMM -#include // system with dot +#include // system with dot #include // STL #include ## Namespace Usage diff --git a/include/cucascade/memory/common.hpp b/include/cucascade/memory/common.hpp index 9ec69be..46a7068 100644 --- a/include/cucascade/memory/common.hpp +++ b/include/cucascade/memory/common.hpp @@ -21,7 +21,11 @@ #include #include +#if __has_include() +#include +#else #include +#endif #include #include diff --git a/include/cucascade/memory/null_device_memory_resource.hpp b/include/cucascade/memory/null_device_memory_resource.hpp index d6cabfa..aab89a8 100644 --- a/include/cucascade/memory/null_device_memory_resource.hpp +++ b/include/cucascade/memory/null_device_memory_resource.hpp @@ -18,7 +18,11 @@ #pragma once #include +#if __has_include() +#include +#else #include +#endif #include diff --git a/include/cucascade/memory/numa_region_pinned_host_allocator.hpp b/include/cucascade/memory/numa_region_pinned_host_allocator.hpp index f772f44..4c55ab3 100644 --- a/include/cucascade/memory/numa_region_pinned_host_allocator.hpp +++ b/include/cucascade/memory/numa_region_pinned_host_allocator.hpp @@ -18,7 +18,11 @@ #pragma once #include +#if __has_include() +#include +#else #include +#endif #include diff --git a/include/cucascade/memory/small_pinned_host_memory_resource.hpp b/include/cucascade/memory/small_pinned_host_memory_resource.hpp index c4497c7..31beede 100644 --- a/include/cucascade/memory/small_pinned_host_memory_resource.hpp +++ b/include/cucascade/memory/small_pinned_host_memory_resource.hpp @@ -20,7 +20,11 @@ #include #include +#if __has_include() +#include +#else #include +#endif #include #include @@ -35,40 +39,19 @@ namespace cucascade { namespace memory { /** - * @brief A pinned host memory allocator combining small slab pools with a bucketed reuse cache for - * large allocations. + * @brief Provides pooled pinned host memory for small and large allocations * - * Requests up to MAX_SLAB_SIZE are served from slab pools (SLAB_SIZES) of pinned host memory. Each - * pool is populated on demand by acquiring one upstream block from the provided - * fixed_size_host_memory_resource and carving it into slabs of the appropriate size. + * Requests up to `MAX_SLAB_SIZE` use slabs carved from the upstream resource. Larger requests whose + * power-of-two bucket fits the configured cache limit can reuse released `cudaHostAlloc` buffers. + * Requests whose bucket exceeds the limit allocate exactly the requested size and bypass the cache. * - * Requests above MAX_SLAB_SIZE are rounded up to a power-of-two bucket (MIN_LARGE_BUCKET at - * minimum) and served from per-bucket free lists of previously released `cudaHostAlloc` buffers; - * misses allocate a fresh buffer. Released buffers are cached rather than freed, up to a - * configurable byte limit, sparing callers such as cuDF's parquet reader the synchronous - * `cudaHostAlloc` / `cudaFreeHost` cost on every request. Reuse in both the slab and large paths is - * ordered by CUDA events recorded on the freeing stream, so a buffer is never handed out while - * another stream may still be using it. - * - * Satisfies the ::cuda::mr::device_accessible and ::cuda::mr::host_accessible - * properties, making it compatible with rmm::host_device_async_resource_ref - * and suitable for use as cuDF's default pinned memory resource. - * - * Typical use: - * @code - * small_pinned_host_memory_resource slab_mr(host_fixed_mr); - * cudf::set_pinned_memory_resource(slab_mr); - * cudf::set_allocate_host_as_pinned_threshold( - * small_pinned_host_memory_resource::MAX_SLAB_SIZE); - * @endcode - * - * This eliminates the pageable H2D transfers that cuDF would otherwise issue - * when building column_device_view metadata arrays for cudf::concatenate. + * Reuse is ordered after work submitted to the stream passed to `deallocate`. The resource + * satisfies the ::cuda::mr::device_accessible and ::cuda::mr::host_accessible properties required + * by rmm::host_device_async_resource_ref. */ class small_pinned_host_memory_resource { public: - /// Maximum allocation size handled by the slab pools. Larger requests are served from the - /// large-allocation cache. + /// Largest request served by a slab pool. static constexpr std::size_t MAX_SLAB_SIZE = 8192; /// Smallest power-of-two bucket used for allocations above MAX_SLAB_SIZE. @@ -78,11 +61,11 @@ class small_pinned_host_memory_resource { static constexpr std::size_t DEFAULT_LARGE_CACHE_LIMIT = 256ull << 20; /** - * @brief Construct with the upstream fixed-size host memory resource. + * @brief Constructs a pinned host memory resource * * @param upstream Block allocator backed by pinned host memory. Must outlive this object. - * @param large_cache_limit_bytes Maximum total bytes, measured in bucket sizes, retained by the - * large-allocation cache. + * @param large_cache_limit_bytes Maximum total bucket capacity retained for large allocations. A + * value below `MIN_LARGE_BUCKET` disables large-allocation caching. */ explicit small_pinned_host_memory_resource( fixed_size_host_memory_resource& upstream, @@ -93,45 +76,58 @@ class small_pinned_host_memory_resource { small_pinned_host_memory_resource(small_pinned_host_memory_resource&&) = delete; small_pinned_host_memory_resource& operator=(small_pinned_host_memory_resource&&) = delete; + /** + * @brief Releases retained buffers and CUDA events + * + * All work using allocations returned by this resource must be complete, and no calls may be in + * flight when it is destroyed. + */ ~small_pinned_host_memory_resource(); /** - * @brief Allocate pinned memory. + * @brief Allocates pinned host memory * - * For @p bytes <= MAX_SLAB_SIZE: rounds up to the next slab boundary (512 / 1 KB / 2 KB / 4 KB / - * 8 KB) and returns a pointer from the matching free list, expanding the pool from upstream if - * the list is empty. + * Requests up to `MAX_SLAB_SIZE` are rounded to the next slab size. A larger request is rounded + * to its power-of-two bucket only when that bucket fits the cache limit; otherwise the resource + * allocates exactly @p bytes. If a direct pinned allocation fails, retained large buffers are + * released and the allocation is retried once. * - * For @p bytes > MAX_SLAB_SIZE: rounds up to the power-of-two bucket and, when the bucket fits - * within the cache limit, returns a cached buffer when the bucket's free list has one, making @p - * stream wait on the buffer's ready event first. On a miss (or a never-cacheable size, which - * skips the cache lookup), allocates a fresh buffer of large_allocation_size(bytes) with - * `cudaHostAlloc(Portable | Mapped)`; if that fails, purges the entire large cache and retries - * once before throwing std::bad_alloc. + * @throw std::bad_alloc If a direct pinned allocation still fails after retained buffers are + * released + * + * @param stream CUDA stream on which reuse dependencies are inserted + * @param bytes Number of bytes requested + * @param alignment Requested alignment; currently not used to select storage + * @return Pointer to at least @p bytes bytes of pinned memory, or `nullptr` when @p bytes is zero */ void* allocate(::cuda::stream_ref stream, std::size_t bytes, std::size_t alignment = alignof(std::max_align_t)); /** - * @brief Return memory to the appropriate free list. + * @brief Deallocates pinned host memory * - * Slabs (@p bytes <= MAX_SLAB_SIZE) are returned to the slab free list. Larger buffers are cached - * in their bucket's free list, evicting the oldest cached entries (in insertion order across - * buckets) when the total would exceed the cache limit; a buffer whose bucket alone exceeds the - * limit is freed with `cudaFreeHost` instead. Both paths record an event on @p stream so reuse - * waits for pending work on the buffer. When no event can be recorded, @p stream is synchronized - * instead (best effort); a large buffer is then freed rather than cached and a slab is cached - * carrying no pending work, so a cached entry never carries pending work that reuse cannot order - * against. + * Cacheable buffers are retained for reuse, while large buffers whose bucket exceeds the cache + * limit are released. The caller must order every prior access to @p ptr before or on @p stream. * - * @p bytes must equal the value passed to the corresponding allocate. + * @param stream CUDA stream ordered after the final use of @p ptr + * @param ptr Pointer returned by this resource, or `nullptr` for a no-op + * @param bytes Original requested allocation size; must match the corresponding call to + * `allocate` + * @param alignment Original requested alignment */ void deallocate(::cuda::stream_ref stream, void* ptr, std::size_t bytes, std::size_t alignment = alignof(std::max_align_t)) noexcept; + /** + * @brief Allocates pinned host memory and synchronizes the default stream + * + * @param bytes Number of bytes requested + * @param alignment Requested alignment; currently not used to select storage + * @return Pointer to at least @p bytes bytes of pinned memory, or `nullptr` when @p bytes is zero + */ void* allocate_sync(std::size_t bytes, std::size_t alignment = alignof(std::max_align_t)) { auto* ptr = allocate(cuda::stream_ref{cudaStream_t{nullptr}}, bytes, alignment); @@ -139,6 +135,13 @@ class small_pinned_host_memory_resource { return ptr; } + /** + * @brief Deallocates pinned host memory and synchronizes the default stream + * + * @param ptr Pointer returned by this resource, or `nullptr` for a no-op + * @param bytes Original requested allocation size + * @param alignment Original requested alignment + */ void deallocate_sync(void* ptr, std::size_t bytes, std::size_t alignment = alignof(std::max_align_t)) noexcept @@ -147,13 +150,27 @@ class small_pinned_host_memory_resource { rmm::cuda_stream_default.synchronize_no_throw(); } - /// Total bytes currently held in the large-allocation cache, measured in bucket sizes. + /** + * @brief Returns the large-buffer capacity currently accounted to the cache + * + * This excludes slab storage, live allocations, and entries already removed for eviction or + * purge. + * + * @return Sum of bucket sizes for cached entries available for reuse + */ [[nodiscard]] std::size_t large_cache_bytes() const; + /** + * @brief Compares memory resource identity + * + * @param other Resource to compare + * @return `true` if @p other is this resource + */ bool operator==(small_pinned_host_memory_resource const& other) const noexcept; /** - * @brief Declares that memory allocated here is accessible from GPU devices. + * @brief Enables the ::cuda::mr::device_accessible property + * * Required to satisfy rmm::host_device_async_resource_ref. */ friend void get_property(small_pinned_host_memory_resource const&, @@ -162,7 +179,8 @@ class small_pinned_host_memory_resource { } /** - * @brief Declares that memory allocated here is accessible from the host. + * @brief Enables the ::cuda::mr::host_accessible property + * * Required to satisfy rmm::host_device_async_resource_ref. */ friend void get_property(small_pinned_host_memory_resource const&, @@ -173,104 +191,86 @@ class small_pinned_host_memory_resource { /// Slab sizes in ascending order. static constexpr std::array SLAB_SIZES{512, 1024, 2048, 4096, 8192}; - /// Returns the index into SLAB_SIZES of the smallest slab >= bytes. + /** + * @brief Finds the smallest slab that can satisfy a request + * + * @param bytes Request size no greater than `MAX_SLAB_SIZE` + * @return Index of the matching entry in `SLAB_SIZES` + */ static std::size_t slab_index_for(std::size_t bytes) noexcept; - /// Returns the bucket for a request above MAX_SLAB_SIZE: the smallest power of two >= bytes, no - /// smaller than MIN_LARGE_BUCKET. The bucket is used for cache keying and eviction accounting; - /// the physical size is large_allocation_size(bytes), equal to the bucket only for cacheable - /// requests. allocate and deallocate both derive the bucket from the request size, so the pairing - /// is deterministic. + /** + * @brief Computes the cache bucket for a large request + * + * @param bytes Request size greater than `MAX_SLAB_SIZE` + * @return Smallest representable power of two no less than @p bytes and `MIN_LARGE_BUCKET`, or @p + * bytes when the next power of two is not representable + */ static std::size_t large_bucket_size_for(std::size_t bytes) noexcept; - /// Returns the physical size cudaHostAlloc is asked for on a cache miss: the bucket when it fits - /// within the cache limit, otherwise exactly @p bytes. A never-cacheable allocation gains nothing - /// from bucket rounding and must not overshoot pinned memory (up to 2x for sizes just past a - /// bucket boundary). Reads only the immutable cache limit, so no lock is needed. + /** + * @brief Computes the storage size for a direct large allocation + * + * @param bytes Requested allocation size + * @return Cache bucket size when it fits the limit, or @p bytes otherwise + */ [[nodiscard]] std::size_t large_allocation_size(std::size_t bytes) const noexcept; - /// Populate the free list for slab @p idx by acquiring one upstream block. - /// Must be called with mutex_ held. + /// Populates slab pool @p slab_idx from an upstream block. Must hold @c mutex_. void expand_pool_locked(std::size_t slab_idx); - /// A CUDA event paired with the device whose context owns it. cudaEventRecord requires the event - /// and the stream to share a CUDA context, so pooled events are segregated by device and an event - /// is only ever recorded on a stream of its own device. @c device is meaningful only when @c - /// handle is non-null. + /// CUDA event and the device on which it was created. Recording succeeds only on a stream + /// associated with the same device. struct device_event { cudaEvent_t handle{nullptr}; int device{-1}; }; - /// A free slab plus, when it was just deallocated, a CUDA event recorded on - /// the freeing stream. Reusing the slab must wait on this event so an - /// in-flight async H2D copy that still reads the slab (e.g. cuDF's parquet - /// stats min/max buffers) completes before another stream overwrites it. - /// A null @c ready.handle means the slab carries no pending work: it was - /// freshly carved and never used, or the freeing stream was synchronized - /// before the slab was cached. + /// Slab available for reuse and an optional event recording its previous use. struct free_slab { void* ptr; device_event ready; }; - /// A released large buffer held for reuse. Like @c free_slab, @c ready captures the freeing - /// stream's last use of the buffer; unlike a slab, a large buffer is never cached without a - /// recorded event (deallocate frees it instead), so @c ready.handle is non-null for every cached - /// entry. @c sequence orders entries across buckets so eviction can drop the oldest first. + /// Large buffer available for reuse, its required dependency, and its insertion order. struct large_cache_entry { void* ptr; device_event ready; std::uint64_t sequence; }; - /// Borrow a timing-disabled CUDA event owned by the calling thread's current device (recycled - /// from @c event_pools_ or newly created). Returns a null handle when the device query or event - /// creation fails; the caller then falls back to freeing (large path) or synchronizing (slab - /// path) instead of caching with ordering. Must hold @c mutex_. + /// Borrows a timing-disabled event created on the current device. Returns an empty event if the + /// device query or event creation fails. Must hold @c mutex_. device_event acquire_event_locked(); - /// Return an event to its device's pool in @c event_pools_ for reuse. Must hold @c mutex_. + /// Returns an event to its device's pool in @c event_pools_ for reuse. Must hold @c mutex_. void release_event_locked(device_event event) noexcept; - /// Pop the oldest cached buffer of @p bucket and make @p stream wait on its ready event, - /// recycling the event. Returns null when the bucket has no cached entries. Must hold @c mutex_. + /// Removes the oldest entry in @p bucket and orders @p stream after its ready event. Returns + /// `nullptr` when the bucket is empty. Must hold @c mutex_. void* try_take_cached_large_locked(std::size_t bucket, ::cuda::stream_ref stream); - /// Remove the oldest cached large buffer across all buckets from the bookkeeping and return it - /// with its ready event still attached, for the caller to pass to sync_and_free_large_victim - /// outside the lock and then recycle the event. Returns a null-ptr entry when the cache is empty. - /// Must hold @c mutex_. + /// Removes and returns the oldest entry across all buckets. Returns an entry with a null pointer + /// when the cache is empty. Must hold @c mutex_. large_cache_entry evict_oldest_large_locked() noexcept; - /// Remove every cached large buffer from the bookkeeping and return them with their ready events - /// still attached, for the caller to pass to sync_and_free_large_victim outside the lock and then - /// recycle the events. Must hold @c mutex_. + /// Removes and returns every cached large buffer. Must hold @c mutex_. std::vector purge_large_cache_locked(); - /// Wait for a victim's recorded work and free its buffer. The ready event may have been recorded - /// on any device's stream, so the buffer is not unpinned until the event has completed - /// (best-effort: a failed or null event skips the wait). Must be called without @c mutex_ held; - /// the caller recycles the event afterwards. + /// Waits for @p victim's ready event, when present, and frees its buffer. If the wait fails, + /// freeing proceeds as a best-effort fallback. Must be called without @c mutex_ held. static void sync_and_free_large_victim(large_cache_entry const& victim) noexcept; fixed_size_host_memory_resource& upstream_; mutable std::mutex mutex_; std::array, 5> free_lists_{}; - /// Idle recycled events keyed by the device that owns them. Events are created in the calling - /// thread's current context, so acquire_event_locked keys by cudaGetDevice() at acquire time; - /// this avoids cudaSetDevice churn, and in deployment the calling thread's current device is its - /// stream's device (the NUMA dispatcher routes by cudaGetDevice()). Should a caller ever pass a - /// stream of another device, the record fails and the safe fallback engages, so correctness never - /// depends on the key, only cache-hit rate. + // Recycled events grouped by the device on which they were created. std::map> event_pools_; std::vector owned_allocations_; - /// Large-allocation cache: per-bucket free lists keyed by bucket size, the current total in - /// bucket bytes, the retention cap, and a monotonic counter stamping insertion order for - /// eviction. + // Large-allocation cache and its bucket-capacity accounting. std::map> large_cache_; std::size_t large_cache_bytes_ = 0; std::size_t const large_cache_limit_bytes_; diff --git a/src/memory/small_pinned_host_memory_resource.cpp b/src/memory/small_pinned_host_memory_resource.cpp index b90b306..737d812 100644 --- a/src/memory/small_pinned_host_memory_resource.cpp +++ b/src/memory/small_pinned_host_memory_resource.cpp @@ -38,14 +38,8 @@ small_pinned_host_memory_resource::small_pinned_host_memory_resource( small_pinned_host_memory_resource::~small_pinned_host_memory_resource() { - // owned_allocations_ destructor returns upstream blocks to the free list. - // free_lists_ entries are raw pointers into those blocks; no individual cleanup needed. - // Cached large buffers come straight from cudaHostAlloc, so each is freed here. - // Destroy every CUDA event we own — both idle (pooled) and still attached to a - // free slab or cached large buffer that was never re-allocated. The events may - // be owned by multiple devices' contexts; cudaEventDestroy carries no - // same-device precondition, so destruction order and the current device are - // irrelevant. + // Slabs are suballocations released by owned_allocations_; only their CUDA events need explicit + // cleanup. for (auto& pool : event_pools_) { for (auto& event : pool.second) { if (event != nullptr) { CUCASCADE_ASSERT_CUDA_SUCCESS(::cudaEventDestroy(event)); } @@ -58,9 +52,7 @@ small_pinned_host_memory_resource::~small_pinned_host_memory_resource() } } } - // Destruction assumes no further allocate/deallocate calls, but work already recorded on a - // cached entry's ready event may still be draining, so each buffer goes through the same - // wait-then-free as eviction and purge before its event is destroyed. + // Cached large buffers own their storage, so wait for recorded work before unpinning them. for (auto& bucket : large_cache_) { for (auto& entry : bucket.second) { sync_and_free_large_victim(entry); @@ -76,8 +68,7 @@ small_pinned_host_memory_resource::acquire_event_locked() { int device = -1; if (::cudaGetDevice(&device) != cudaSuccess) { - // Best effort: a null handle routes the caller to its no-event fallback, and allocation must - // never throw here. Clear the sticky error and carry on. + // Clear CUDA's sticky error and return an empty event for the caller's no-event fallback. (void)::cudaGetLastError(); return {}; } @@ -90,8 +81,7 @@ small_pinned_host_memory_resource::acquire_event_locked() cudaEvent_t event = nullptr; // Timing is not needed; disabling it makes record/wait cheaper. if (::cudaEventCreateWithFlags(&event, cudaEventDisableTiming) != cudaSuccess) { - // Best effort, as above: the caller handles the null handle; allocation must never throw - // here. Clear the sticky error and carry on. + // Clear CUDA's sticky error and return an empty event for the caller's no-event fallback. (void)::cudaGetLastError(); return {}; } @@ -111,15 +101,8 @@ void* small_pinned_host_memory_resource::allocate(cuda::stream_ref stream, [[maybe_unused]] std::size_t alignment) { if (bytes == 0) { return nullptr; } - // cuDF calls get_pinned_memory_resource() directly from some code paths (e.g. join/sort - // staging buffers) that bypass the allocate_host_as_pinned threshold check. Serve those - // with cudaHostAlloc(Portable) so the memory remains pinned AND DMA-accessible from - // every CUDA context (multi-GPU consumers need the Portable flag; cudaMallocHost / - // cudaHostAllocDefault produce memory that is only DMA-accessible from the allocating - // device's context, which under CUDA 13+ makes cudaMemcpyBatchAsync reject cross-device - // sources with cudaErrorInvalidValue). cuDF 26.04+ may access hostdevice_vector memory - // directly from GPU kernels (e.g. detect_malformed_pages), so returning pageable memory - // here would cause cudaErrorIllegalAddress. + // Large requests bypass upstream slabs but remain portable, mapped pinned memory for + // cross-context use. if (bytes > MAX_SLAB_SIZE) { std::size_t const bucket = large_bucket_size_for(bytes); bool const cacheable = bucket <= large_cache_limit_bytes_; @@ -129,17 +112,13 @@ void* small_pinned_host_memory_resource::allocate(cuda::stream_ref stream, if (void* cached = try_take_cached_large_locked(bucket, stream)) { return cached; } } void* ptr = nullptr; - // Portable + Mapped — see numa_region_pinned_host_allocator.cpp comment. - auto err = ::cudaHostAlloc(&ptr, alloc_bytes, cudaHostAllocPortable | cudaHostAllocMapped); + auto err = ::cudaHostAlloc(&ptr, alloc_bytes, cudaHostAllocPortable | cudaHostAllocMapped); if (err == cudaSuccess) { return ptr; } // Clear the sticky error so a successful retry does not leave cudaGetLastError consumers // seeing a stale allocation failure. (void)::cudaGetLastError(); - // Cached buffers are the only pinned memory this class can give back under pressure: - // release them all and retry once (even a never-cacheable request benefits, since purging - // frees the pinned memory its retry needs). Victims are synchronized and freed outside the - // lock so slab traffic does not stall behind the waits; their events are recycled in one - // batch under the re-taken lock. + // Release retained pinned memory before retrying. Wait and free outside the mutex because + // event synchronization may block. std::vector purged; { std::lock_guard lock(mutex_); @@ -171,12 +150,8 @@ void* small_pinned_host_memory_resource::allocate(cuda::stream_ref stream, if (free_lists_[idx].empty()) { expand_pool_locked(idx); } free_slab slab = free_lists_[idx].back(); free_lists_[idx].pop_back(); - // If this slab was recently deallocated, its ready event captures the freeing - // stream's last use (e.g. an in-flight async H2D copy still reading the slab). - // Make the reusing stream wait for it so we cannot overwrite the slab before - // that copy completes. Recording the event again (on a later deallocate) does - // not disturb this already-enqueued wait, so the event is safe to recycle. A - // null handle means the slab carries no pending work, so no wait is needed. + // cudaStreamWaitEvent snapshots the event's current record, so the handle can be recycled after + // the wait is enqueued. if (slab.ready.handle != nullptr) { CUCASCADE_ASSERT_CUDA_SUCCESS(::cudaStreamWaitEvent(stream.get(), slab.ready.handle, 0)); release_event_locked(slab.ready); @@ -193,24 +168,18 @@ void small_pinned_host_memory_resource::deallocate(cuda::stream_ref stream, if (bytes > MAX_SLAB_SIZE) { std::size_t const bucket = large_bucket_size_for(bytes); if (bucket > large_cache_limit_bytes_) { - // Too big to ever cache: free directly. Any pending work on the buffer was issued by the - // caller on its own device, which is the case cudaFreeHost's implicit synchronization - // covers; no ready event has been recorded for this buffer. + // This bucket cannot fit under the retention cap, so bypass the cache. CUCASCADE_ASSERT_CUDA_SUCCESS(::cudaFreeHost(ptr)); return; } - // Evict until the incoming bucket fits, then cache it. A victim's pending work may live on - // any device's stream, so each iteration unhooks one victim under the lock, waits on its - // ready event and frees it outside the lock (slab traffic must not stall behind the wait), - // and then recycles the event. + // Remove victims under the mutex, then wait and free outside it because event synchronization + // may block. while (true) { large_cache_entry victim{nullptr, {}, 0}; bool free_instead_of_caching = false; { std::lock_guard lock(mutex_); if (large_cache_bytes_ + bucket <= large_cache_limit_bytes_) { - // As in the slab path below, the recorded event defers reuse until pending work on - // the freeing stream completes. device_event event = acquire_event_locked(); if (event.handle != nullptr && ::cudaEventRecord(event.handle, stream.get()) != cudaSuccess) { @@ -223,21 +192,15 @@ void small_pinned_host_memory_resource::deallocate(cuda::stream_ref stream, large_cache_bytes_ += bucket; return; } - // Reuse ordering is carried solely by the recorded event, so a buffer whose event - // could not be recorded is freed (outside the lock, below) instead of cached. This - // branch is rare because pooled events are segregated by device: it fires only on - // event-creation failure or a stream that does not belong to the caller's current - // device. + // Cache only after a successful record; otherwise reuse could race pending work. free_instead_of_caching = true; } else { victim = evict_oldest_large_locked(); } } if (free_instead_of_caching) { - // With no event recorded on the freeing stream, ordering comes from draining the stream - // before the free, exactly as in the slab fallback below. The stream may belong to - // another device's context, which cudaFreeHost's implicit synchronization is not - // documented to cover. Best effort on failure: clear the sticky error and free anyway. + // Without a recorded event, drain the deallocation stream before freeing. Deallocation is + // noexcept, so clear synchronization errors and continue as a best-effort fallback. if (::cudaStreamSynchronize(stream.get()) != cudaSuccess) { (void)::cudaGetLastError(); } CUCASCADE_ASSERT_CUDA_SUCCESS(::cudaFreeHost(ptr)); return; @@ -259,8 +222,7 @@ void small_pinned_host_memory_resource::deallocate(cuda::stream_ref stream, std::size_t idx = slab_index_for(bytes); { std::lock_guard lock(mutex_); - // Record an event on the freeing stream so a future reuse of this slab can wait - // for any still-pending work on it (the async H2D copy in cuDF's stats filter). + // Record the final stream use so a later allocation can insert a dependency. device_event event = acquire_event_locked(); if (event.handle != nullptr) { if (::cudaEventRecord(event.handle, stream.get()) == cudaSuccess) { @@ -271,10 +233,8 @@ void small_pinned_host_memory_resource::deallocate(cuda::stream_ref stream, release_event_locked(event); } } - // No event could be recorded, so drain the freeing stream before caching: a null ready handle - // promises the slab carries no pending work, and synchronizing here is what keeps that promise. - // Best effort on failure: a stream broken enough to fail synchronize is a context where ordering - // is already lost, and leaking the slab would be the only alternative. + // Without a recorded event, synchronize before publishing the slab. Deallocation is noexcept, so + // clear synchronization errors and continue as a best-effort fallback. if (::cudaStreamSynchronize(stream.get()) != cudaSuccess) { (void)::cudaGetLastError(); } // The push must happen after the synchronize completes; pushing first would let a racing // allocate hand the slab out before the freeing stream drains. @@ -299,8 +259,8 @@ std::size_t small_pinned_host_memory_resource::slab_index_for(std::size_t bytes) std::size_t small_pinned_host_memory_resource::large_bucket_size_for(std::size_t bytes) noexcept { if (bytes <= MIN_LARGE_BUCKET) { return MIN_LARGE_BUCKET; } - // bit_ceil is undefined when the next power of two is unrepresentable. Such a request cannot - // be satisfied anyway, so pass it through unrounded and let cudaHostAlloc reject it. + // std::bit_ceil is undefined when its result is not representable; use the exact request as the + // bucket in that case. constexpr std::size_t max_bucket = std::size_t{1} << (std::numeric_limits::digits - 1); if (bytes > max_bucket) { return bytes; } @@ -329,8 +289,8 @@ void* small_pinned_host_memory_resource::try_take_cached_large_locked(std::size_ it->second.pop_front(); if (it->second.empty()) { large_cache_.erase(it); } large_cache_bytes_ -= bucket; - // Same reuse discipline as the slab path: make the reusing stream wait for the freeing - // stream's last use of this buffer before it can be overwritten. + // cudaStreamWaitEvent snapshots the event's current record, so the handle can be recycled after + // the wait is enqueued. if (entry.ready.handle != nullptr) { CUCASCADE_ASSERT_CUDA_SUCCESS(::cudaStreamWaitEvent(stream.get(), entry.ready.handle, 0)); release_event_locked(entry.ready); @@ -381,9 +341,8 @@ small_pinned_host_memory_resource::purge_large_cache_locked() void small_pinned_host_memory_resource::sync_and_free_large_victim( large_cache_entry const& victim) noexcept { - // The ready event may target any device's stream; wait for it before unpinning the pages so - // an in-flight DMA cannot read freed memory. Best effort: on a failed wait, clear the sticky - // error and free anyway. + // Wait for recorded use before unpinning. Deallocation is noexcept, so clear event errors and + // free as a best-effort fallback. if (victim.ready.handle != nullptr && ::cudaEventSynchronize(victim.ready.handle) != cudaSuccess) { (void)::cudaGetLastError(); @@ -393,7 +352,6 @@ void small_pinned_host_memory_resource::sync_and_free_large_victim( void small_pinned_host_memory_resource::expand_pool_locked(std::size_t slab_idx) { - // Acquire one upstream block and carve it into slabs. std::size_t upstream_block_size = upstream_.get_block_size(); auto allocation = upstream_.allocate_multiple_blocks(upstream_block_size); diff --git a/test/memory/test_reservation_aware_resource_adaptor.cpp b/test/memory/test_reservation_aware_resource_adaptor.cpp index 3195af8..4e3a711 100644 --- a/test/memory/test_reservation_aware_resource_adaptor.cpp +++ b/test/memory/test_reservation_aware_resource_adaptor.cpp @@ -39,7 +39,11 @@ #include #include +#if __has_include() +#include +#else #include +#endif #include #include diff --git a/test/memory/test_small_pinned_host_memory_resource.cpp b/test/memory/test_small_pinned_host_memory_resource.cpp index 9ff8143..75aa941 100644 --- a/test/memory/test_small_pinned_host_memory_resource.cpp +++ b/test/memory/test_small_pinned_host_memory_resource.cpp @@ -256,7 +256,6 @@ TEST_CASE("Large allocations do not interfere with slab pool", "[small_pinned]") constexpr std::size_t big_size = 16384; constexpr std::size_t small_size = 512; - // Allocate a large chunk (goes to malloc) auto* big = f.slab_mr.allocate(rmm::cuda_stream_view{}, big_size); REQUIRE(big != nullptr); From 5df1ea0e7ac78dea0713e3a580207cb0286b1740 Mon Sep 17 00:00:00 2001 From: Kevin Kristensen Date: Sat, 15 Aug 2026 19:35:18 +0000 Subject: [PATCH 3/3] docs(memory): revert unrelated allocator docs --- .../small_pinned_host_memory_resource.hpp | 33 ++----------------- 1 file changed, 3 insertions(+), 30 deletions(-) diff --git a/include/cucascade/memory/small_pinned_host_memory_resource.hpp b/include/cucascade/memory/small_pinned_host_memory_resource.hpp index 31beede..47cade2 100644 --- a/include/cucascade/memory/small_pinned_host_memory_resource.hpp +++ b/include/cucascade/memory/small_pinned_host_memory_resource.hpp @@ -121,13 +121,6 @@ class small_pinned_host_memory_resource { std::size_t bytes, std::size_t alignment = alignof(std::max_align_t)) noexcept; - /** - * @brief Allocates pinned host memory and synchronizes the default stream - * - * @param bytes Number of bytes requested - * @param alignment Requested alignment; currently not used to select storage - * @return Pointer to at least @p bytes bytes of pinned memory, or `nullptr` when @p bytes is zero - */ void* allocate_sync(std::size_t bytes, std::size_t alignment = alignof(std::max_align_t)) { auto* ptr = allocate(cuda::stream_ref{cudaStream_t{nullptr}}, bytes, alignment); @@ -135,13 +128,6 @@ class small_pinned_host_memory_resource { return ptr; } - /** - * @brief Deallocates pinned host memory and synchronizes the default stream - * - * @param ptr Pointer returned by this resource, or `nullptr` for a no-op - * @param bytes Original requested allocation size - * @param alignment Original requested alignment - */ void deallocate_sync(void* ptr, std::size_t bytes, std::size_t alignment = alignof(std::max_align_t)) noexcept @@ -160,17 +146,10 @@ class small_pinned_host_memory_resource { */ [[nodiscard]] std::size_t large_cache_bytes() const; - /** - * @brief Compares memory resource identity - * - * @param other Resource to compare - * @return `true` if @p other is this resource - */ bool operator==(small_pinned_host_memory_resource const& other) const noexcept; /** - * @brief Enables the ::cuda::mr::device_accessible property - * + * @brief Declares that memory allocated here is accessible from GPU devices. * Required to satisfy rmm::host_device_async_resource_ref. */ friend void get_property(small_pinned_host_memory_resource const&, @@ -179,8 +158,7 @@ class small_pinned_host_memory_resource { } /** - * @brief Enables the ::cuda::mr::host_accessible property - * + * @brief Declares that memory allocated here is accessible from the host. * Required to satisfy rmm::host_device_async_resource_ref. */ friend void get_property(small_pinned_host_memory_resource const&, @@ -191,12 +169,7 @@ class small_pinned_host_memory_resource { /// Slab sizes in ascending order. static constexpr std::array SLAB_SIZES{512, 1024, 2048, 4096, 8192}; - /** - * @brief Finds the smallest slab that can satisfy a request - * - * @param bytes Request size no greater than `MAX_SLAB_SIZE` - * @return Index of the matching entry in `SLAB_SIZES` - */ + /// Returns the index into SLAB_SIZES of the smallest slab >= bytes. static std::size_t slab_index_for(std::size_t bytes) noexcept; /**