From 8ebdb6ccb6b67ae7c62be1d07bcffebf91672dd6 Mon Sep 17 00:00:00 2001 From: David Wendt Date: Thu, 6 Aug 2026 15:44:34 -0400 Subject: [PATCH 1/2] Make internal host_allocator stream ordered --- .../cudf/detail/utilities/host_memory.hpp | 39 +++- .../cudf/detail/utilities/host_vector.hpp | 3 - .../detail/utilities/vector_factories.hpp | 10 +- cpp/src/utilities/host_memory.cpp | 205 +++++++++--------- 4 files changed, 151 insertions(+), 106 deletions(-) diff --git a/cpp/include/cudf/detail/utilities/host_memory.hpp b/cpp/include/cudf/detail/utilities/host_memory.hpp index d12b7c71226e..5fd92cb075ab 100644 --- a/cpp/include/cudf/detail/utilities/host_memory.hpp +++ b/cpp/include/cudf/detail/utilities/host_memory.hpp @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2024, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2024-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * SPDX-License-Identifier: Apache-2.0 */ @@ -13,6 +13,30 @@ #include namespace cudf::detail { + +/** + * @brief Property tag for the CUDA native host memory pool. + * + * Advertised by `cuda_host_pinned_pool_memory_resource` (the pool backed by + * `cudaMallocFromPoolAsync` / `cudaFreeAsync`). The typed resource ref below carries it through + * type erasure so callers can select the right pool without going through the public + * `get_pinned_memory_resource()` API. + */ +struct stream_ordered_host_accessible_t {}; + +/** + * @brief Typed resource ref for the stream-ordered pinned host pool. + * + * Carrying `stream_ordered_host_accessible_t` in the template parameters means callers that + * construct `rmm_host_allocator` from this ref get a concrete type that still carries the + * property (before it is erased into `rmm::host_async_resource_ref`). This is the return type + * of `get_stream_ordered_pinned_memory_resource()`. + */ +using stream_ordered_host_device_async_resource_ref = + cuda::mr::resource_ref; + /** * @brief Get the memory resource to be used for pageable memory allocations. * @@ -20,6 +44,17 @@ namespace cudf::detail { */ CUDF_EXPORT rmm::host_async_resource_ref get_pageable_memory_resource(); +/** + * @brief Get the stream-ordered pinned memory resource. + * + * The underlying resource uses `cudaMallocFromPoolAsync` / `cudaFreeAsync` so both allocation and + * deallocation are truly stream-ordered. + * + * @return A `stream_ordered_host_device_async_resource_ref` backed by the default pinned pool + */ +CUDF_EXPORT stream_ordered_host_device_async_resource_ref +get_stream_ordered_pinned_memory_resource(); + /** * @brief Get the allocator to be used for the host memory allocation. * @@ -31,7 +66,7 @@ template rmm_host_allocator get_host_allocator(std::size_t size, rmm::cuda_stream_view stream) { if (size * sizeof(T) <= get_allocate_host_as_pinned_threshold()) { - return {get_pinned_memory_resource(), stream}; + return {get_stream_ordered_pinned_memory_resource(), stream}; } return {get_pageable_memory_resource(), stream}; } diff --git a/cpp/include/cudf/detail/utilities/host_vector.hpp b/cpp/include/cudf/detail/utilities/host_vector.hpp index a65a1879ac5a..7771dc24a984 100644 --- a/cpp/include/cudf/detail/utilities/host_vector.hpp +++ b/cpp/include/cudf/detail/utilities/host_vector.hpp @@ -141,9 +141,6 @@ class rmm_host_allocator { { if (cnt > this->max_size()) { throw std::bad_alloc(); } // end if auto const result = mr.allocate(stream, cnt * sizeof(value_type), alignof(value_type)); - // Synchronize to ensure the memory is allocated before thrust::host_vector initialization - // TODO: replace thrust::host_vector with a type that does not require synchronization - stream.synchronize(); return static_cast(result); } diff --git a/cpp/include/cudf/detail/utilities/vector_factories.hpp b/cpp/include/cudf/detail/utilities/vector_factories.hpp index f82ba7aaece2..7e0136970766 100644 --- a/cpp/include/cudf/detail/utilities/vector_factories.hpp +++ b/cpp/include/cudf/detail/utilities/vector_factories.hpp @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2021-2026, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2021-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * SPDX-License-Identifier: Apache-2.0 */ @@ -494,7 +494,8 @@ host_vector make_host_vector(Container const& c, template host_vector make_empty_pinned_vector(size_t capacity, rmm::cuda_stream_view stream) { - auto result = host_vector({cudf::get_pinned_memory_resource(), stream}); + auto result = + host_vector(rmm_host_allocator{get_stream_ordered_pinned_memory_resource(), stream}); result.reserve(capacity); return result; } @@ -502,7 +503,7 @@ host_vector make_empty_pinned_vector(size_t capacity, rmm::cuda_stream_view s /** * @brief Asynchronously construct a pinned `cudf::detail::host_vector` of the given size * - * @note This function may not synchronize `stream` after the copy. + * @note This function does not synchronize `stream`. * * @tparam T The type of the vector data * @param size The number of elements in the created vector @@ -512,7 +513,8 @@ host_vector make_empty_pinned_vector(size_t capacity, rmm::cuda_stream_view s template host_vector make_pinned_vector_async(size_t size, rmm::cuda_stream_view stream) { - return host_vector(size, {cudf::get_pinned_memory_resource(), stream}); + return host_vector(size, + rmm_host_allocator{get_stream_ordered_pinned_memory_resource(), stream}); } /** diff --git a/cpp/src/utilities/host_memory.cpp b/cpp/src/utilities/host_memory.cpp index 3d6e56471c50..b4167bc36483 100644 --- a/cpp/src/utilities/host_memory.cpp +++ b/cpp/src/utilities/host_memory.cpp @@ -4,6 +4,7 @@ */ #include +#include #include #include #include @@ -12,8 +13,6 @@ #include #include -#include -#include #include #include @@ -23,8 +22,6 @@ #include #include #include -#include -#include namespace cudf { @@ -75,51 +72,78 @@ void aligned_host_deallocate(void* ptr, } } -class pinned_pool_with_fallback_memory_resource { - using upstream_mr = rmm::mr::pinned_host_memory_resource; - using host_pooled_mr = rmm::mr::pool_memory_resource; - - struct fallback_state { - mutable std::shared_mutex mutex; - std::unordered_set allocations; - }; - - private: - upstream_mr upstream_mr_{}; - size_t initial_pool_size_{0}; - size_t max_pool_size_{0}; - // Raw pointer to avoid a segfault when the pool is destroyed on exit - host_pooled_mr* pool_{nullptr}; - cuda::stream_ref stream_{cudf::detail::global_cuda_stream_pool().get_stream().value()}; - - // Wrapped in shared_ptr so the outer class is copyable (required by any_resource) - std::shared_ptr fallback_{std::make_shared()}; - +/** + * @brief Stream-ordered pinned host memory pool backed by the CUDA native memory pool API. + * + * Uses `cudaMallocFromPoolAsync` / `cudaFreeAsync` with a pool created for + * `cudaMemAllocationTypePinned + cudaMemLocationTypeHost`. This provides two benefits over the + * old `rmm::mr::pool_memory_resource` approach: + * + * 1. **No sync on allocation.** The CUDA runtime guarantees that recycled blocks have no + * outstanding GPU accesses, so `rmm_host_allocator::allocate()` can skip + * `stream.synchronize()`. + * 2. **Stream-ordered free.** Callers can destroy a `host_vector` immediately after enqueuing + * a device-to-host copy; the runtime defers the actual reclaim until the stream reaches the + * free point. + * + * The pool grows on demand. `release_threshold` controls how many bytes the pool retains before + * returning memory to the OS (mirrors the old max-pool-size semantics). + */ +class cuda_host_pinned_pool_memory_resource { public: - pinned_pool_with_fallback_memory_resource(size_t initial_size, size_t max_size) - : // rmm requires the pool size to be a multiple of 256 bytes - initial_pool_size_{rmm::align_up(initial_size, rmm::CUDA_ALLOCATION_ALIGNMENT)}, - max_pool_size_{rmm::align_up(max_size, rmm::CUDA_ALLOCATION_ALIGNMENT)}, - pool_{new host_pooled_mr(upstream_mr_, initial_pool_size_, max_pool_size_)} + explicit cuda_host_pinned_pool_memory_resource(std::size_t release_threshold) + { + cudaMemPoolProps props{}; + props.allocType = cudaMemAllocationTypePinned; + props.handleTypes = cudaMemHandleTypeNone; + props.location.type = cudaMemLocationTypeHost; + props.location.id = 0; // id is ignored for cudaMemLocationTypeHost + CUDF_CUDA_TRY(cudaMemPoolCreate(&pool_, &props)); + + // Keep up to release_threshold bytes in the pool before releasing to the OS. + CUDF_CUDA_TRY( + cudaMemPoolSetAttribute(pool_, cudaMemPoolAttrReleaseThreshold, &release_threshold)); + + // Unlike cudaMallocHost, host pools are not device-accessible by default. + int current_device = 0; + CUDF_CUDA_TRY(cudaGetDevice(¤t_device)); + cudaMemAccessDesc access{}; + access.location.type = cudaMemLocationTypeDevice; + access.location.id = current_device; + access.flags = cudaMemAccessFlagsProtReadWrite; + CUDF_CUDA_TRY(cudaMemPoolSetAccess(pool_, &access, 1)); + + CUDF_LOG_INFO("CUDA host pinned pool created, release threshold = %zu bytes, device = %d", + release_threshold, + current_device); + } + + // The pool handle is an opaque value, so copies are cheap. Copies share the same underlying + // pool; operator== distinguishes them. We intentionally do NOT destroy the pool in the + // destructor: cuda::mr::resource_ref stores a copy of the resource via type erasure, and + // calling cudaMemPoolDestroy at process exit can race with CUDA teardown (same reasoning as + // the raw pool_ pointer in the old pinned_pool_with_fallback_memory_resource). + cuda_host_pinned_pool_memory_resource(cuda_host_pinned_pool_memory_resource const&) = default; + cuda_host_pinned_pool_memory_resource& operator=(cuda_host_pinned_pool_memory_resource const&) = + default; + + // clang-tidy will complain about these get_property friends because they are completely + // unused at runtime and only exist for tag introspection by CCCL, so we ignore linting. + // This masks a real issue if we ever want to compile with clang, though, which is that the + // function will actually be compiled out by clang. The same goes for the other get_property + // definitions in this file. + friend void get_property(cuda_host_pinned_pool_memory_resource const&, // NOLINT + cuda::mr::host_accessible) noexcept { - CUDF_LOG_INFO( - "Pinned pool initial size = %zu, max size = %zu", initial_pool_size_, max_pool_size_); } - // clang-tidy will complain about this function because it is completely - // unused at runtime and only exist for tag introspection by CCCL, so we - // ignore linting. This masks a real issue if we ever want to compile with - // clang, though, which is that the function will actually be compiled out by - // clang. If cudf were ever to try to support clang as a compile we would - // need to force the compiler to emit this symbol. The same goes for the - // other get_property definitions in this file. - friend void get_property(pinned_pool_with_fallback_memory_resource const&, // NOLINT + friend void get_property(cuda_host_pinned_pool_memory_resource const&, // NOLINT cuda::mr::device_accessible) noexcept { } - friend void get_property(pinned_pool_with_fallback_memory_resource const&, // NOLINT - cuda::mr::host_accessible) noexcept + friend void get_property(cuda_host_pinned_pool_memory_resource const&, // NOLINT + cudf::detail::stream_ordered_host_accessible_t) noexcept { } @@ -137,74 +161,47 @@ class pinned_pool_with_fallback_memory_resource { void* allocate(cuda::stream_ref stream, std::size_t bytes, - std::size_t alignment = rmm::CUDA_ALLOCATION_ALIGNMENT) + [[maybe_unused]] std::size_t alignment = rmm::CUDA_ALLOCATION_ALIGNMENT) { - if (max_pool_size_ == 0) { return upstream_mr_.allocate(stream, bytes, alignment); } - - try { - return pool_->allocate(stream, bytes, alignment); - } catch (...) { - CUDF_LOG_INFO("Pinned pool exhausted, falling back to new pinned allocation for %zu bytes", - bytes); - // fall back to upstream - auto* ptr = upstream_mr_.allocate(stream, bytes, alignment); - - { - std::unique_lock lock(fallback_->mutex); - fallback_->allocations.insert(ptr); - } - - return ptr; - } + // cudaMallocFromPoolAsync guarantees at least 256-byte alignment; cudf never requests more. + void* ptr = nullptr; + auto const err = cudaMallocFromPoolAsync(&ptr, bytes, pool_, stream.get()); + if (err != cudaSuccess) { throw std::bad_alloc(); } + return ptr; } void deallocate(cuda::stream_ref stream, void* ptr, - std::size_t bytes, - std::size_t alignment = rmm::CUDA_ALLOCATION_ALIGNMENT) noexcept + [[maybe_unused]] std::size_t bytes, + [[maybe_unused]] std::size_t alignment = rmm::CUDA_ALLOCATION_ALIGNMENT) noexcept { - if (max_pool_size_ == 0) { - upstream_mr_.deallocate(stream, ptr, bytes, alignment); - return; - } - - bool is_fallback{false}; - { - std::shared_lock lock(fallback_->mutex); - is_fallback = fallback_->allocations.find(ptr) != fallback_->allocations.end(); - } - - if (is_fallback) { - { - std::unique_lock lock(fallback_->mutex); - fallback_->allocations.erase(ptr); - } - upstream_mr_.deallocate(stream, ptr, bytes, alignment); - } else { - pool_->deallocate(stream, ptr, bytes, alignment); - } + [[maybe_unused]] auto err = cudaFreeAsync(ptr, stream.get()); } - bool operator==(pinned_pool_with_fallback_memory_resource const& other) const noexcept + bool operator==(cuda_host_pinned_pool_memory_resource const& other) const noexcept { - return pool_ == other.pool_ && stream_ == other.stream_; + return pool_ == other.pool_; } - bool operator!=(pinned_pool_with_fallback_memory_resource const& other) const noexcept + bool operator!=(cuda_host_pinned_pool_memory_resource const& other) const noexcept { return !(*this == other); } + + private: + cudaMemPool_t pool_{}; }; -static_assert(cuda::mr::resource_with, - "Pinned pool mr must be accessible from both host and device"); + cuda::mr::host_accessible, + cudf::detail::stream_ordered_host_accessible_t>, + "CUDA host pinned pool mr must be host/device accessible and stream-ordered"); -CUDF_EXPORT rmm::host_device_async_resource_ref& make_default_pinned_mr( +CUDF_EXPORT cuda_host_pinned_pool_memory_resource& make_default_pinned_mr( std::optional config_size) { - static pinned_pool_with_fallback_memory_resource mr = [config_size]() { + static cuda_host_pinned_pool_memory_resource mr = [config_size]() { auto const initial_size = [&config_size]() -> size_t { if (auto const env_val = getenv("LIBCUDF_PINNED_POOL_SIZE"); env_val != nullptr) { return std::atol(env_val); @@ -213,22 +210,21 @@ CUDF_EXPORT rmm::host_device_async_resource_ref& make_default_pinned_mr( if (config_size.has_value()) { return *config_size; } auto const total = rmm::available_device_memory().second; - // 0.5% of the total device memory, capped at 64MB + // 0.5% of total device memory, capped at 64 MB return std::min(total / 200, size_t{64} * 1024 * 1024); }(); - auto const max_size = [&initial_size]() -> size_t { + auto const release_threshold = [&initial_size]() -> size_t { if (auto const env_val = getenv("LIBCUDF_PINNED_POOL_MAX_SIZE"); env_val != nullptr) { return std::atol(env_val); } return initial_size * 16; }(); - return pinned_pool_with_fallback_memory_resource{initial_size, max_size}; + return cuda_host_pinned_pool_memory_resource{release_threshold}; }(); - static rmm::host_device_async_resource_ref mr_ref{mr}; - return mr_ref; + return mr; } CUDF_EXPORT std::mutex& host_mr_mutex() @@ -245,11 +241,11 @@ CUDF_EXPORT rmm::host_device_async_resource_ref& make_host_mr( bool configured = false; if (mr_ref == nullptr) { configured = true; - mr_ref = &make_default_pinned_mr(opts ? opts->pool_size : std::nullopt); + auto& pool = make_default_pinned_mr(opts ? opts->pool_size : std::nullopt); + static rmm::host_device_async_resource_ref pool_ref{pool}; + mr_ref = &pool_ref; } - // If the user passed an out param to detect whether this call configured a resource - // set the result if (did_configure != nullptr) { *did_configure = configured; } return *mr_ref; @@ -262,6 +258,15 @@ CUDF_EXPORT rmm::host_device_async_resource_ref& host_mr() return mr_ref; } +// Returns a typed ref that carries stream_ordered_host_accessible_t — used internally so that +// rmm_host_allocator can skip stream.synchronize() in allocate(). +CUDF_EXPORT cudf::detail::stream_ordered_host_device_async_resource_ref& stream_ordered_host_mr() +{ + static cudf::detail::stream_ordered_host_device_async_resource_ref mr_ref = + make_default_pinned_mr(std::nullopt); + return mr_ref; +} + class new_delete_memory_resource { public: void* allocate_sync(std::size_t bytes, std::size_t alignment = rmm::CUDA_ALLOCATION_ALIGNMENT) @@ -371,6 +376,12 @@ CUDF_EXPORT rmm::host_async_resource_ref get_pageable_memory_resource() return mr_ref; } +CUDF_EXPORT stream_ordered_host_device_async_resource_ref +get_stream_ordered_pinned_memory_resource() +{ + return stream_ordered_host_mr(); +} + } // namespace detail } // namespace cudf From 1d5299029d6445fd1e155dfffd27cf36aa3334fe Mon Sep 17 00:00:00 2001 From: David Wendt Date: Mon, 10 Aug 2026 16:40:43 -0400 Subject: [PATCH 2/2] copy device itr with use_cpu_for_cluster_computation --- cpp/include/cudf_test/tdigest_utilities.hpp | 8 ++-- .../quantiles/tdigest/tdigest_aggregation.cu | 47 +++++++++---------- 2 files changed, 26 insertions(+), 29 deletions(-) diff --git a/cpp/include/cudf_test/tdigest_utilities.hpp b/cpp/include/cudf_test/tdigest_utilities.hpp index dc252d14b08a..e6013c9abf9c 100644 --- a/cpp/include/cudf_test/tdigest_utilities.hpp +++ b/cpp/include/cudf_test/tdigest_utilities.hpp @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2022-2026, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * SPDX-License-Identifier: Apache-2.0 */ @@ -128,10 +128,8 @@ void tdigest_minmax_compare(cudf::tdigest::tdigest_column_view const& tdv, auto expected_max = static_cast(max_scalar->value()); double tdv_min, tdv_max; - EXPECT_EQ(cudaMemcpy(&tdv_min, tdv.min_begin(), sizeof(double), cudaMemcpyDeviceToHost), - cudaSuccess); - EXPECT_EQ(cudaMemcpy(&tdv_max, tdv.max_begin(), sizeof(double), cudaMemcpyDeviceToHost), - cudaSuccess); + EXPECT_EQ(cudaMemcpy(&tdv_min, tdv.min_begin(), sizeof(double), cudaMemcpyDefault), cudaSuccess); + EXPECT_EQ(cudaMemcpy(&tdv_max, tdv.max_begin(), sizeof(double), cudaMemcpyDefault), cudaSuccess); EXPECT_EQ(tdv_min, expected_min); EXPECT_EQ(tdv_max, expected_max); diff --git a/cpp/src/quantiles/tdigest/tdigest_aggregation.cu b/cpp/src/quantiles/tdigest/tdigest_aggregation.cu index ffa3a3d7a28a..ee36cb4625cd 100644 --- a/cpp/src/quantiles/tdigest/tdigest_aggregation.cu +++ b/cpp/src/quantiles/tdigest/tdigest_aggregation.cu @@ -1536,43 +1536,42 @@ std::unique_ptr merge_tdigests(tdigest_column_view const& tdv, // if we will be at least partially using the CPU here, move the important values into pinned // and reference those instead. if (use_cpu_for_cluster_computation(num_groups)) { - auto pinned_mr = cudf::get_pinned_memory_resource(); - - rmm::device_uvector _p_group_offsets(num_groups + 1, stream, pinned_mr); + rmm::device_uvector d_p_group_offsets(num_groups + 1, stream); thrust::copy(rmm::exec_policy_nosync(stream, cudf::get_current_device_resource_ref()), group_offsets, - group_offsets + _p_group_offsets.size(), - _p_group_offsets.begin()); - cudf::device_span p_group_offsets(_p_group_offsets); - - rmm::device_uvector p_cumulative_weights(cumulative_weights, stream, pinned_mr); - - rmm::device_uvector p_tdigest_offsets(tdigest_offsets.size(), stream, pinned_mr); - thrust::copy(rmm::exec_policy_nosync(stream, cudf::get_current_device_resource_ref()), - tdigest_offsets.begin(), - tdigest_offsets.begin() + p_tdigest_offsets.size(), - p_tdigest_offsets.begin()); - - rmm::device_uvector _p_group_labels(num_group_labels, stream, pinned_mr); + group_offsets + d_p_group_offsets.size(), + d_p_group_offsets.begin()); + auto _p_group_offsets = cudf::detail::make_pinned_vector_async(d_p_group_offsets, stream); + cudf::device_span p_group_offsets(_p_group_offsets.data(), + _p_group_offsets.size()); + + auto _p_tdigest_offsets = cudf::detail::make_pinned_vector_async( + device_span(tdigest_offsets.data(), tdigest_offsets.size()), + stream); + auto p_tdigest_offsets = + cuda::std::span{_p_tdigest_offsets.data(), _p_tdigest_offsets.size()}; + + rmm::device_uvector d_p_group_labels(num_group_labels, stream); thrust::copy(rmm::exec_policy_nosync(stream, cudf::get_current_device_resource_ref()), group_labels, group_labels + num_group_labels, - _p_group_labels.begin()); - cudf::device_span p_group_labels(_p_group_labels); + d_p_group_labels.begin()); + auto _p_group_labels = cudf::detail::make_pinned_vector_async(d_p_group_labels, stream); + cudf::device_span p_group_labels(_p_group_labels.data(), + _p_group_labels.size()); stream.synchronize(); + auto pinned_mr = cudf::get_pinned_memory_resource(); + rmm::device_uvector p_cumulative_weights(cumulative_weights, stream, pinned_mr); return generate_group_cluster_info( delta, num_groups, nearest_value_centroid_weights{ - p_cumulative_weights.begin(), p_group_offsets, p_tdigest_offsets.begin()}, + p_cumulative_weights.begin(), p_group_offsets, p_tdigest_offsets.data()}, centroid_group_info{ - p_cumulative_weights.begin(), p_group_offsets, p_tdigest_offsets.begin()}, + p_cumulative_weights.begin(), p_group_offsets, p_tdigest_offsets.data()}, cumulative_centroid_weight{ - p_cumulative_weights.begin(), - p_group_labels, - p_group_offsets, - cuda::std::span{p_tdigest_offsets.begin(), p_tdigest_offsets.size()}}, + p_cumulative_weights.begin(), p_group_labels, p_group_offsets, p_tdigest_offsets}, has_nulls, stream, mr);