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 cd2748f..47cade2 100644 --- a/include/cucascade/memory/small_pinned_host_memory_resource.hpp +++ b/include/cucascade/memory/small_pinned_host_memory_resource.hpp @@ -20,11 +20,18 @@ #include #include +#if __has_include() +#include +#else #include +#endif #include #include #include +#include +#include +#include #include #include @@ -32,68 +39,82 @@ namespace cucascade { namespace memory { /** - * @brief A small slab allocator backed by fixed_size_host_memory_resource. + * @brief Provides pooled pinned host memory for small and 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` 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. * - * 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. - /// Requests larger than this use pageable memory. + /// 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. + 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. + * @brief Constructs a pinned 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 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); + 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; 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: falls back to cudaMallocHost (pinned). + * @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 pool. + * @brief Deallocates pinned host memory + * + * 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. * - * 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. + * @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, @@ -115,6 +136,16 @@ class small_pinned_host_memory_resource { rmm::cuda_stream_default.synchronize_no_throw(); } + /** + * @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; + bool operator==(small_pinned_host_memory_resource const& other) const noexcept; /** @@ -141,32 +172,82 @@ 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; - /// Populate the free list for slab @p idx by acquiring one upstream block. - /// Must be called with mutex_ held. + /** + * @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; + + /** + * @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; + + /// Populates slab pool @p slab_idx from an upstream block. Must hold @c mutex_. void expand_pool_locked(std::size_t slab_idx); - /// 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. + /// 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}; + }; + + /// Slab available for reuse and an optional event recording its previous use. 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(); + /// 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; + }; - /// Return an event to @c event_pool_ for reuse. Must hold @c mutex_. - void release_event_locked(cudaEvent_t event) noexcept; + /// 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(); + + /// 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; + + /// 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); + + /// 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; + + /// Removes and returns every cached large buffer. Must hold @c mutex_. + std::vector purge_large_cache_locked(); + + /// 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_{}; - std::vector event_pool_; + + // Recycled events grouped by the device on which they were created. + std::map> event_pools_; + std::vector owned_allocations_; + + // 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_; + 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) { } 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. - // 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)); } + // 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)); } + } } 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)); + } + } + } + // 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); + 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) { + // Clear CUDA's sticky error and return an empty event for the caller's no-event fallback. + (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. + // Clear CUDA's sticky error and return an empty event for the caller's no-event fallback. (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) { 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_; + 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(); + // 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_); + 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 +150,96 @@ 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 - // 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); + // 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); } 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_) { + // This bucket cannot fit under the retention cap, so bypass the cache. + CUCASCADE_ASSERT_CUDA_SUCCESS(::cudaFreeHost(ptr)); + return; + } + // 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_) { + 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; + } + // 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) { + // 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; + } + 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 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) { + 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}); + // 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. + std::lock_guard lock(mutex_); + free_lists_[idx].push_back(free_slab{ptr, {}}); } bool small_pinned_host_memory_resource::operator==( @@ -150,9 +256,102 @@ 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; } + // 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; } + 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; + // 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); + } + 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 +{ + // 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(); + } + 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. std::size_t upstream_block_size = upstream_.get_block_size(); auto allocation = upstream_.allocate_multiple_blocks(upstream_block_size); @@ -161,7 +360,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_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 349adba..75aa941 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; @@ -254,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); @@ -272,3 +273,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"); +}