From 163ad977e923ddb20a3747ce957f44cb81b5cb67 Mon Sep 17 00:00:00 2001 From: Vukasin Milovanovic Date: Mon, 6 May 2024 09:53:16 -0700 Subject: [PATCH 01/27] pool with fallback --- cpp/src/io/utilities/config_utils.cpp | 132 ++++++++++++++++++++++++-- 1 file changed, 126 insertions(+), 6 deletions(-) diff --git a/cpp/src/io/utilities/config_utils.cpp b/cpp/src/io/utilities/config_utils.cpp index 2f7a6131e3d..e54896fe650 100644 --- a/cpp/src/io/utilities/config_utils.cpp +++ b/cpp/src/io/utilities/config_utils.cpp @@ -16,10 +16,12 @@ #include "config_utils.hpp" +#include #include #include #include +#include #include #include @@ -87,16 +89,134 @@ bool is_stable_enabled() { return is_all_enabled() or get_env_policy() == usage_ } // namespace nvcomp_integration -inline std::mutex& host_mr_lock() +using rmm_pinned_pool_t = rmm::mr::pool_memory_resource; + +class fixed_pinned_pool_memory_resource { + private: + rmm_pinned_pool_t* _pool; + void* pool_begin_{}; + void* pool_end_{}; + size_t pool_size_{}; + cuda::stream_ref stream_{cudf::get_default_stream().value()}; + rmm::mr::pinned_host_memory_resource fallback_mr_{}; + + public: + fixed_pinned_pool_memory_resource(rmm_pinned_pool_t* pool) + : _pool(pool), pool_size_{pool->pool_size()} + { + // allocate from the pinned pool the full size to figure out + // our beginning and end address. + if (pool_size_ != 0) { + pool_begin_ = pool->allocate_async(pool_size_, stream_); + pool_end_ = static_cast(static_cast(pool_begin_) + pool_size_); + pool->deallocate_async(pool_begin_, pool_size_, stream_); + } + } + + void* do_allocate_async(std::size_t bytes, std::size_t alignment, cuda::stream_ref stream) + { + if (bytes <= pool_size_) { + try { + return _pool->allocate_async(bytes, alignment, stream); + } catch (const std::exception& unused) { + } + } + + // std::cout << "Falling back!\n"; + return fallback_mr_.allocate_async(bytes, alignment, stream); + } + void do_deallocate_async(void* ptr, + std::size_t bytes, + std::size_t alignment, + cuda::stream_ref stream) noexcept + { + if (bytes <= pool_size_ && ptr >= pool_begin_ && ptr <= pool_end_) { + _pool->deallocate_async(ptr, bytes, alignment, stream); + } else { + fallback_mr_.deallocate_async(ptr, bytes, alignment, stream); + } + } + + void* allocate_async(std::size_t bytes, cuda::stream_ref stream) + { + return do_allocate_async(bytes, rmm::RMM_DEFAULT_HOST_ALIGNMENT, stream); + } + + void* allocate_async(std::size_t bytes, std::size_t alignment, cuda::stream_ref stream) + { + return do_allocate_async(bytes, alignment, stream); + } + + void* allocate(std::size_t bytes, + [[maybe_unused]] std::size_t alignment = rmm::RMM_DEFAULT_HOST_ALIGNMENT) + { + return allocate_async(bytes, stream_); + } + + void deallocate_async(void* ptr, std::size_t bytes, cuda::stream_ref stream) noexcept + { + return do_deallocate_async(ptr, bytes, rmm::RMM_DEFAULT_HOST_ALIGNMENT, stream); + } + + void deallocate_async(void* ptr, + std::size_t bytes, + std::size_t alignment, + cuda::stream_ref stream) noexcept + { + return do_deallocate_async(ptr, bytes, alignment, stream); + } + + void deallocate(void* ptr, + std::size_t bytes, + std::size_t alignment = rmm::RMM_DEFAULT_HOST_ALIGNMENT) noexcept + { + return deallocate_async(ptr, bytes, alignment, stream_); + } + + bool operator==(const fixed_pinned_pool_memory_resource&) const { return true; } + + bool operator!=(const fixed_pinned_pool_memory_resource&) const { return false; } + + friend void get_property(fixed_pinned_pool_memory_resource const&, + cuda::mr::device_accessible) noexcept + { + } + + friend void get_property(fixed_pinned_pool_memory_resource const&, + cuda::mr::host_accessible) noexcept + { + } +}; + +inline rmm::host_async_resource_ref default_pinned_mr() { - static std::mutex map_lock; - return map_lock; + using host_pooled_mr = rmm::mr::pool_memory_resource; + + auto const size = []() -> size_t { + if (auto const env_val = getenv("LIBCUDF_PINNED_POOL_SIZE")) { return std::atol(env_val); } + + size_t free{}, total{}; + cudaMemGetInfo(&free, &total); + // 0.5% of the total device memory, capped at 100MB + return std::min((total / 200 + 255) & ~255, size_t{100} * 1024 * 1024); + }(); + + CUDF_LOG_INFO("Pinned pool size = {}", size); + + // make the pool with max size equal to the initial size + static std::shared_ptr pool_mr = std::make_shared( + std::make_shared().get(), size, size); + + static std::shared_ptr mr = + std::make_shared(pool_mr.get()); + + return *mr; } -inline rmm::host_async_resource_ref default_pinned_mr() +inline std::mutex& host_mr_lock() { - static rmm::mr::pinned_host_memory_resource default_mr{}; - return default_mr; + static std::mutex map_lock; + return map_lock; } CUDF_EXPORT inline auto& host_mr() From 1e850d6dfb497961fd312a7a863ef26414781630 Mon Sep 17 00:00:00 2001 From: Vukasin Milovanovic Date: Mon, 6 May 2024 14:51:47 -0700 Subject: [PATCH 02/27] don't use default pool --- .../cudf/detail/utilities/stream_pool.hpp | 5 +++++ cpp/src/io/utilities/config_utils.cpp | 17 +++++++++++++---- 2 files changed, 18 insertions(+), 4 deletions(-) diff --git a/cpp/include/cudf/detail/utilities/stream_pool.hpp b/cpp/include/cudf/detail/utilities/stream_pool.hpp index 19ef26a10cb..23b54387755 100644 --- a/cpp/include/cudf/detail/utilities/stream_pool.hpp +++ b/cpp/include/cudf/detail/utilities/stream_pool.hpp @@ -81,6 +81,11 @@ class cuda_stream_pool { */ cuda_stream_pool* create_global_cuda_stream_pool(); +/** + * @brief Get the global stream pool. + */ +cuda_stream_pool& global_cuda_stream_pool(); + /** * @brief Acquire a set of `cuda_stream_view` objects and synchronize them to an event on another * stream. diff --git a/cpp/src/io/utilities/config_utils.cpp b/cpp/src/io/utilities/config_utils.cpp index e54896fe650..c9d0056c261 100644 --- a/cpp/src/io/utilities/config_utils.cpp +++ b/cpp/src/io/utilities/config_utils.cpp @@ -16,6 +16,7 @@ #include "config_utils.hpp" +#include #include #include #include @@ -97,7 +98,7 @@ class fixed_pinned_pool_memory_resource { void* pool_begin_{}; void* pool_end_{}; size_t pool_size_{}; - cuda::stream_ref stream_{cudf::get_default_stream().value()}; + cuda::stream_ref stream_{cudf::detail::global_cuda_stream_pool().get_stream(0).value()}; rmm::mr::pinned_host_memory_resource fallback_mr_{}; public: @@ -122,7 +123,6 @@ class fixed_pinned_pool_memory_resource { } } - // std::cout << "Falling back!\n"; return fallback_mr_.allocate_async(bytes, alignment, stream); } void do_deallocate_async(void* ptr, @@ -150,7 +150,9 @@ class fixed_pinned_pool_memory_resource { void* allocate(std::size_t bytes, [[maybe_unused]] std::size_t alignment = rmm::RMM_DEFAULT_HOST_ALIGNMENT) { - return allocate_async(bytes, stream_); + auto const result = allocate_async(bytes, stream_); + stream_.wait(); + return result; } void deallocate_async(void* ptr, std::size_t bytes, cuda::stream_ref stream) noexcept @@ -170,7 +172,8 @@ class fixed_pinned_pool_memory_resource { std::size_t bytes, std::size_t alignment = rmm::RMM_DEFAULT_HOST_ALIGNMENT) noexcept { - return deallocate_async(ptr, bytes, alignment, stream_); + deallocate_async(ptr, bytes, alignment, stream_); + stream_.wait(); } bool operator==(const fixed_pinned_pool_memory_resource&) const { return true; } @@ -219,6 +222,12 @@ inline std::mutex& host_mr_lock() return map_lock; } +inline rmm::host_async_resource_ref old_default_pinned_mr() +{ + static rmm::mr::pinned_host_memory_resource default_mr{}; + return default_mr; +} + CUDF_EXPORT inline auto& host_mr() { static rmm::host_async_resource_ref host_mr = default_pinned_mr(); From 3be42ba43fb342c64ae9151873d39ad1c2f1c43a Mon Sep 17 00:00:00 2001 From: Vukasin Milovanovic Date: Mon, 6 May 2024 14:53:49 -0700 Subject: [PATCH 03/27] fix allocator copy assignment --- cpp/include/cudf/detail/utilities/rmm_host_vector.hpp | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/cpp/include/cudf/detail/utilities/rmm_host_vector.hpp b/cpp/include/cudf/detail/utilities/rmm_host_vector.hpp index 858501877b0..4d05ed049bc 100644 --- a/cpp/include/cudf/detail/utilities/rmm_host_vector.hpp +++ b/cpp/include/cudf/detail/utilities/rmm_host_vector.hpp @@ -124,7 +124,8 @@ class rmm_host_allocator { */ rmm_host_allocator& operator=(rmm_host_allocator const& other) { - mr = other.mr; + mr = other.mr; + stream = other.stream; return *this; } From 395dcf17a5b29b5dd9f6b66ecb60acd0f7bf0caf Mon Sep 17 00:00:00 2001 From: Vukasin Milovanovic Date: Mon, 6 May 2024 15:27:26 -0700 Subject: [PATCH 04/27] fix ver2 --- .../cudf/detail/utilities/rmm_host_vector.hpp | 25 ------------------- 1 file changed, 25 deletions(-) diff --git a/cpp/include/cudf/detail/utilities/rmm_host_vector.hpp b/cpp/include/cudf/detail/utilities/rmm_host_vector.hpp index 4d05ed049bc..46280188561 100644 --- a/cpp/include/cudf/detail/utilities/rmm_host_vector.hpp +++ b/cpp/include/cudf/detail/utilities/rmm_host_vector.hpp @@ -109,31 +109,6 @@ class rmm_host_allocator { { } - /** - * @brief Copy constructor - */ - rmm_host_allocator(rmm_host_allocator const& other) = default; - - /** - * @brief Move constructor - */ - rmm_host_allocator(rmm_host_allocator&& other) = default; - - /** - * @brief Assignment operator - */ - rmm_host_allocator& operator=(rmm_host_allocator const& other) - { - mr = other.mr; - stream = other.stream; - return *this; - } - - /** - * @brief rmm_host_allocator's null destructor does nothing. - */ - inline ~rmm_host_allocator() {} - /** * @brief This method allocates storage for objects in host memory. * From 0873b1f35932d90508ca09cd907d159ba8503b9e Mon Sep 17 00:00:00 2001 From: Vukasin Milovanovic Date: Mon, 6 May 2024 15:51:14 -0700 Subject: [PATCH 05/27] copyright --- cpp/include/cudf/detail/utilities/stream_pool.hpp | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/cpp/include/cudf/detail/utilities/stream_pool.hpp b/cpp/include/cudf/detail/utilities/stream_pool.hpp index 23b54387755..e19cc3ec2f7 100644 --- a/cpp/include/cudf/detail/utilities/stream_pool.hpp +++ b/cpp/include/cudf/detail/utilities/stream_pool.hpp @@ -1,5 +1,5 @@ /* - * Copyright (c) 2023, NVIDIA CORPORATION. + * Copyright (c) 2023-2024, NVIDIA CORPORATION. * * Licensed under the Apache License, Version 2.0 (the "License"); * you may not use this file except in compliance with the License. From ff18a2106bbc733d674c78c3d73e49e3a81e20bb Mon Sep 17 00:00:00 2001 From: Vukasin Milovanovic Date: Tue, 7 May 2024 11:57:32 -0700 Subject: [PATCH 06/27] fix operator== --- cpp/include/cudf/detail/utilities/rmm_host_vector.hpp | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/cpp/include/cudf/detail/utilities/rmm_host_vector.hpp b/cpp/include/cudf/detail/utilities/rmm_host_vector.hpp index 46280188561..6901a19473e 100644 --- a/cpp/include/cudf/detail/utilities/rmm_host_vector.hpp +++ b/cpp/include/cudf/detail/utilities/rmm_host_vector.hpp @@ -159,7 +159,10 @@ class rmm_host_allocator { * @param x The other \p rmm_host_allocator of interest. * @return This method always returns \c true. */ - inline bool operator==(rmm_host_allocator const& x) const { return x.mr == mr; } + inline bool operator==(rmm_host_allocator const& x) const + { + return x.mr == mr && x.stream == stream; + } /** * @brief This method tests this \p rmm_host_allocator for inequality From 854c0abfd842a8d34bbbb2d8ee2bd4a984b07142 Mon Sep 17 00:00:00 2001 From: Vukasin Milovanovic Date: Tue, 7 May 2024 12:59:31 -0700 Subject: [PATCH 07/27] simplify pool creation --- cpp/src/io/utilities/config_utils.cpp | 35 ++++++++++++--------------- 1 file changed, 16 insertions(+), 19 deletions(-) diff --git a/cpp/src/io/utilities/config_utils.cpp b/cpp/src/io/utilities/config_utils.cpp index c9d0056c261..f2638e692ba 100644 --- a/cpp/src/io/utilities/config_utils.cpp +++ b/cpp/src/io/utilities/config_utils.cpp @@ -93,24 +93,27 @@ bool is_stable_enabled() { return is_all_enabled() or get_env_policy() == usage_ using rmm_pinned_pool_t = rmm::mr::pool_memory_resource; class fixed_pinned_pool_memory_resource { + using upstream_mr = rmm::mr::pinned_host_memory_resource; + using host_pooled_mr = rmm::mr::pool_memory_resource; + private: - rmm_pinned_pool_t* _pool; + upstream_mr upstream_mr_{}; + size_t pool_size_{}; + host_pooled_mr* pool_; void* pool_begin_{}; void* pool_end_{}; - size_t pool_size_{}; cuda::stream_ref stream_{cudf::detail::global_cuda_stream_pool().get_stream(0).value()}; - rmm::mr::pinned_host_memory_resource fallback_mr_{}; public: - fixed_pinned_pool_memory_resource(rmm_pinned_pool_t* pool) - : _pool(pool), pool_size_{pool->pool_size()} + fixed_pinned_pool_memory_resource(size_t size) + : pool_size_{size}, pool_{new host_pooled_mr(upstream_mr_, size, size)} { // allocate from the pinned pool the full size to figure out // our beginning and end address. if (pool_size_ != 0) { - pool_begin_ = pool->allocate_async(pool_size_, stream_); + pool_begin_ = pool_->allocate_async(pool_size_, stream_); pool_end_ = static_cast(static_cast(pool_begin_) + pool_size_); - pool->deallocate_async(pool_begin_, pool_size_, stream_); + pool_->deallocate_async(pool_begin_, pool_size_, stream_); } } @@ -118,12 +121,12 @@ class fixed_pinned_pool_memory_resource { { if (bytes <= pool_size_) { try { - return _pool->allocate_async(bytes, alignment, stream); + return pool_->allocate_async(bytes, alignment, stream); } catch (const std::exception& unused) { } } - return fallback_mr_.allocate_async(bytes, alignment, stream); + return upstream_mr_.allocate_async(bytes, alignment, stream); } void do_deallocate_async(void* ptr, std::size_t bytes, @@ -131,9 +134,9 @@ class fixed_pinned_pool_memory_resource { cuda::stream_ref stream) noexcept { if (bytes <= pool_size_ && ptr >= pool_begin_ && ptr <= pool_end_) { - _pool->deallocate_async(ptr, bytes, alignment, stream); + pool_->deallocate_async(ptr, bytes, alignment, stream); } else { - fallback_mr_.deallocate_async(ptr, bytes, alignment, stream); + upstream_mr_.deallocate_async(ptr, bytes, alignment, stream); } } @@ -193,8 +196,6 @@ class fixed_pinned_pool_memory_resource { inline rmm::host_async_resource_ref default_pinned_mr() { - using host_pooled_mr = rmm::mr::pool_memory_resource; - auto const size = []() -> size_t { if (auto const env_val = getenv("LIBCUDF_PINNED_POOL_SIZE")) { return std::atol(env_val); } @@ -207,13 +208,9 @@ inline rmm::host_async_resource_ref default_pinned_mr() CUDF_LOG_INFO("Pinned pool size = {}", size); // make the pool with max size equal to the initial size - static std::shared_ptr pool_mr = std::make_shared( - std::make_shared().get(), size, size); - - static std::shared_ptr mr = - std::make_shared(pool_mr.get()); + static fixed_pinned_pool_memory_resource mr{size}; - return *mr; + return mr; } inline std::mutex& host_mr_lock() From 5bf0ce4b17af5752016f5a5b5099b4df44db129c Mon Sep 17 00:00:00 2001 From: Vukasin Milovanovic Date: Tue, 7 May 2024 13:25:18 -0700 Subject: [PATCH 08/27] namespace; comments --- cpp/src/io/utilities/config_utils.cpp | 20 ++++++++++---------- 1 file changed, 10 insertions(+), 10 deletions(-) diff --git a/cpp/src/io/utilities/config_utils.cpp b/cpp/src/io/utilities/config_utils.cpp index f2638e692ba..7d9afd7f370 100644 --- a/cpp/src/io/utilities/config_utils.cpp +++ b/cpp/src/io/utilities/config_utils.cpp @@ -90,8 +90,9 @@ bool is_stable_enabled() { return is_all_enabled() or get_env_policy() == usage_ } // namespace nvcomp_integration -using rmm_pinned_pool_t = rmm::mr::pool_memory_resource; +} // namespace detail +namespace { class fixed_pinned_pool_memory_resource { using upstream_mr = rmm::mr::pinned_host_memory_resource; using host_pooled_mr = rmm::mr::pool_memory_resource; @@ -99,6 +100,7 @@ class fixed_pinned_pool_memory_resource { private: upstream_mr upstream_mr_{}; size_t pool_size_{}; + // Raw pointer to avoid a segfault when the pool is destroyed on exit host_pooled_mr* pool_; void* pool_begin_{}; void* pool_end_{}; @@ -108,8 +110,7 @@ class fixed_pinned_pool_memory_resource { fixed_pinned_pool_memory_resource(size_t size) : pool_size_{size}, pool_{new host_pooled_mr(upstream_mr_, size, size)} { - // allocate from the pinned pool the full size to figure out - // our beginning and end address. + // Allocate full size from the pinned pool to figure out the beginning and end address if (pool_size_ != 0) { pool_begin_ = pool_->allocate_async(pool_size_, stream_); pool_end_ = static_cast(static_cast(pool_begin_) + pool_size_); @@ -230,21 +231,20 @@ CUDF_EXPORT inline auto& host_mr() static rmm::host_async_resource_ref host_mr = default_pinned_mr(); return host_mr; } - -} // namespace detail +} // namespace rmm::host_async_resource_ref set_host_memory_resource(rmm::host_async_resource_ref mr) { - std::lock_guard lock{detail::host_mr_lock()}; - auto last_mr = detail::host_mr(); - detail::host_mr() = mr; + std::lock_guard lock{host_mr_lock()}; + auto last_mr = host_mr(); + host_mr() = mr; return last_mr; } rmm::host_async_resource_ref get_host_memory_resource() { - std::lock_guard lock{detail::host_mr_lock()}; - return detail::host_mr(); + std::lock_guard lock{host_mr_lock()}; + return host_mr(); } } // namespace cudf::io From ff4d7f61dd44100737c9f1daca788dc94cb39704 Mon Sep 17 00:00:00 2001 From: Vukasin Milovanovic Date: Tue, 7 May 2024 13:28:16 -0700 Subject: [PATCH 09/27] clean up --- cpp/src/io/utilities/config_utils.cpp | 7 ------- 1 file changed, 7 deletions(-) diff --git a/cpp/src/io/utilities/config_utils.cpp b/cpp/src/io/utilities/config_utils.cpp index 7d9afd7f370..35a7654e73c 100644 --- a/cpp/src/io/utilities/config_utils.cpp +++ b/cpp/src/io/utilities/config_utils.cpp @@ -17,7 +17,6 @@ #include "config_utils.hpp" #include -#include #include #include @@ -220,12 +219,6 @@ inline std::mutex& host_mr_lock() return map_lock; } -inline rmm::host_async_resource_ref old_default_pinned_mr() -{ - static rmm::mr::pinned_host_memory_resource default_mr{}; - return default_mr; -} - CUDF_EXPORT inline auto& host_mr() { static rmm::host_async_resource_ref host_mr = default_pinned_mr(); From 80b5963205ce4dd17e55cbcbcd70f92c6d7e231e Mon Sep 17 00:00:00 2001 From: Vukasin Milovanovic Date: Tue, 7 May 2024 15:43:26 -0700 Subject: [PATCH 10/27] mild polish --- cpp/src/io/utilities/config_utils.cpp | 31 ++++++++++++++++----------- 1 file changed, 18 insertions(+), 13 deletions(-) diff --git a/cpp/src/io/utilities/config_utils.cpp b/cpp/src/io/utilities/config_utils.cpp index 35a7654e73c..6698a587c89 100644 --- a/cpp/src/io/utilities/config_utils.cpp +++ b/cpp/src/io/utilities/config_utils.cpp @@ -98,11 +98,11 @@ class fixed_pinned_pool_memory_resource { private: upstream_mr upstream_mr_{}; - size_t pool_size_{}; + size_t pool_size_{0}; // Raw pointer to avoid a segfault when the pool is destroyed on exit - host_pooled_mr* pool_; - void* pool_begin_{}; - void* pool_end_{}; + host_pooled_mr* pool_{nullptr}; + void* pool_begin_{nullptr}; + void* pool_end_{nullptr}; cuda::stream_ref stream_{cudf::detail::global_cuda_stream_pool().get_stream(0).value()}; public: @@ -150,10 +150,9 @@ class fixed_pinned_pool_memory_resource { return do_allocate_async(bytes, alignment, stream); } - void* allocate(std::size_t bytes, - [[maybe_unused]] std::size_t alignment = rmm::RMM_DEFAULT_HOST_ALIGNMENT) + void* allocate(std::size_t bytes, std::size_t alignment = rmm::RMM_DEFAULT_HOST_ALIGNMENT) { - auto const result = allocate_async(bytes, stream_); + auto const result = do_allocate_async(bytes, alignment, stream_); stream_.wait(); return result; } @@ -179,17 +178,23 @@ class fixed_pinned_pool_memory_resource { stream_.wait(); } - bool operator==(const fixed_pinned_pool_memory_resource&) const { return true; } + bool operator==(fixed_pinned_pool_memory_resource const& other) const + { + return pool_ == other.pool_ and stream_ == other.stream_; + } - bool operator!=(const fixed_pinned_pool_memory_resource&) const { return false; } + bool operator!=(fixed_pinned_pool_memory_resource const& other) const + { + return !operator==(other); + } - friend void get_property(fixed_pinned_pool_memory_resource const&, - cuda::mr::device_accessible) noexcept + [[maybe_unused]] friend void get_property(fixed_pinned_pool_memory_resource const&, + cuda::mr::device_accessible) noexcept { } - friend void get_property(fixed_pinned_pool_memory_resource const&, - cuda::mr::host_accessible) noexcept + [[maybe_unused]] friend void get_property(fixed_pinned_pool_memory_resource const&, + cuda::mr::host_accessible) noexcept { } }; From 0122038e6accf6a999b065530aa1f0fd6269235c Mon Sep 17 00:00:00 2001 From: Vukasin Milovanovic Date: Thu, 9 May 2024 11:03:43 -0700 Subject: [PATCH 11/27] remove inline --- cpp/src/io/utilities/config_utils.cpp | 7 ++++--- 1 file changed, 4 insertions(+), 3 deletions(-) diff --git a/cpp/src/io/utilities/config_utils.cpp b/cpp/src/io/utilities/config_utils.cpp index 6698a587c89..cd005ba38b8 100644 --- a/cpp/src/io/utilities/config_utils.cpp +++ b/cpp/src/io/utilities/config_utils.cpp @@ -199,7 +199,7 @@ class fixed_pinned_pool_memory_resource { } }; -inline rmm::host_async_resource_ref default_pinned_mr() +rmm::host_async_resource_ref default_pinned_mr() { auto const size = []() -> size_t { if (auto const env_val = getenv("LIBCUDF_PINNED_POOL_SIZE")) { return std::atol(env_val); } @@ -218,17 +218,18 @@ inline rmm::host_async_resource_ref default_pinned_mr() return mr; } -inline std::mutex& host_mr_lock() +std::mutex& host_mr_lock() { static std::mutex map_lock; return map_lock; } -CUDF_EXPORT inline auto& host_mr() +CUDF_EXPORT auto& host_mr() { static rmm::host_async_resource_ref host_mr = default_pinned_mr(); return host_mr; } + } // namespace rmm::host_async_resource_ref set_host_memory_resource(rmm::host_async_resource_ref mr) From 60030da003f3664b673df4a636c592f5c3b906cf Mon Sep 17 00:00:00 2001 From: Vukasin Milovanovic Date: Thu, 9 May 2024 11:04:03 -0700 Subject: [PATCH 12/27] scoped_lock --- cpp/src/io/utilities/config_utils.cpp | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/cpp/src/io/utilities/config_utils.cpp b/cpp/src/io/utilities/config_utils.cpp index cd005ba38b8..cc87e07aa86 100644 --- a/cpp/src/io/utilities/config_utils.cpp +++ b/cpp/src/io/utilities/config_utils.cpp @@ -218,7 +218,7 @@ rmm::host_async_resource_ref default_pinned_mr() return mr; } -std::mutex& host_mr_lock() +std::mutex& host_mr_mutex() { static std::mutex map_lock; return map_lock; @@ -234,7 +234,7 @@ CUDF_EXPORT auto& host_mr() rmm::host_async_resource_ref set_host_memory_resource(rmm::host_async_resource_ref mr) { - std::lock_guard lock{host_mr_lock()}; + std::scoped_lock lock{host_mr_mutex()}; auto last_mr = host_mr(); host_mr() = mr; return last_mr; @@ -242,7 +242,7 @@ rmm::host_async_resource_ref set_host_memory_resource(rmm::host_async_resource_r rmm::host_async_resource_ref get_host_memory_resource() { - std::lock_guard lock{host_mr_lock()}; + std::scoped_lock lock{host_mr_mutex()}; return host_mr(); } From a62377e6b5df36c30a1df87914a8c4f1280521ad Mon Sep 17 00:00:00 2001 From: Vukasin Milovanovic Date: Thu, 9 May 2024 11:07:24 -0700 Subject: [PATCH 13/27] try --- cpp/src/io/utilities/config_utils.cpp | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/cpp/src/io/utilities/config_utils.cpp b/cpp/src/io/utilities/config_utils.cpp index cc87e07aa86..1f842ed12a3 100644 --- a/cpp/src/io/utilities/config_utils.cpp +++ b/cpp/src/io/utilities/config_utils.cpp @@ -205,7 +205,7 @@ rmm::host_async_resource_ref default_pinned_mr() if (auto const env_val = getenv("LIBCUDF_PINNED_POOL_SIZE")) { return std::atol(env_val); } size_t free{}, total{}; - cudaMemGetInfo(&free, &total); + CUDF_CUDA_TRY(cudaMemGetInfo(&free, &total)); // 0.5% of the total device memory, capped at 100MB return std::min((total / 200 + 255) & ~255, size_t{100} * 1024 * 1024); }(); From 6733c455ac30dc4b418279d5b7a57223d07d4736 Mon Sep 17 00:00:00 2001 From: Vukasin Milovanovic Date: Thu, 9 May 2024 11:23:09 -0700 Subject: [PATCH 14/27] clean up --- cpp/src/io/utilities/config_utils.cpp | 14 ++++++++------ 1 file changed, 8 insertions(+), 6 deletions(-) diff --git a/cpp/src/io/utilities/config_utils.cpp b/cpp/src/io/utilities/config_utils.cpp index 1f842ed12a3..7b88df1a4e3 100644 --- a/cpp/src/io/utilities/config_utils.cpp +++ b/cpp/src/io/utilities/config_utils.cpp @@ -109,12 +109,12 @@ class fixed_pinned_pool_memory_resource { fixed_pinned_pool_memory_resource(size_t size) : pool_size_{size}, pool_{new host_pooled_mr(upstream_mr_, size, size)} { + if (pool_size_ == 0) { return; } + // Allocate full size from the pinned pool to figure out the beginning and end address - if (pool_size_ != 0) { - pool_begin_ = pool_->allocate_async(pool_size_, stream_); - pool_end_ = static_cast(static_cast(pool_begin_) + pool_size_); - pool_->deallocate_async(pool_begin_, pool_size_, stream_); - } + pool_begin_ = pool_->allocate_async(pool_size_, stream_); + pool_end_ = static_cast(static_cast(pool_begin_) + pool_size_); + pool_->deallocate_async(pool_begin_, pool_size_, stream_); } void* do_allocate_async(std::size_t bytes, std::size_t alignment, cuda::stream_ref stream) @@ -202,7 +202,9 @@ class fixed_pinned_pool_memory_resource { rmm::host_async_resource_ref default_pinned_mr() { auto const size = []() -> size_t { - if (auto const env_val = getenv("LIBCUDF_PINNED_POOL_SIZE")) { return std::atol(env_val); } + if (auto const env_val = getenv("LIBCUDF_PINNED_POOL_SIZE"); env_val != nullptr) { + return std::atol(env_val); + } size_t free{}, total{}; CUDF_CUDA_TRY(cudaMemGetInfo(&free, &total)); From abf40a8f9289a0142efd617ec4dd39ac93424f5c Mon Sep 17 00:00:00 2001 From: Vukasin Milovanovic Date: Thu, 9 May 2024 11:30:04 -0700 Subject: [PATCH 15/27] clarify try-catch fallback --- cpp/src/io/utilities/config_utils.cpp | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/cpp/src/io/utilities/config_utils.cpp b/cpp/src/io/utilities/config_utils.cpp index 7b88df1a4e3..5595b8e2bd1 100644 --- a/cpp/src/io/utilities/config_utils.cpp +++ b/cpp/src/io/utilities/config_utils.cpp @@ -122,7 +122,8 @@ class fixed_pinned_pool_memory_resource { if (bytes <= pool_size_) { try { return pool_->allocate_async(bytes, alignment, stream); - } catch (const std::exception& unused) { + } catch (...) { + // If the pool is exhausted, fall back to the upstream memory resource } } From fa7dce70af6d45790ebfc136850f6772826c3dd9 Mon Sep 17 00:00:00 2001 From: Vukasin Milovanovic Date: Thu, 9 May 2024 11:34:37 -0700 Subject: [PATCH 16/27] remove export --- cpp/src/io/utilities/config_utils.cpp | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/cpp/src/io/utilities/config_utils.cpp b/cpp/src/io/utilities/config_utils.cpp index 5595b8e2bd1..fec5ca88a12 100644 --- a/cpp/src/io/utilities/config_utils.cpp +++ b/cpp/src/io/utilities/config_utils.cpp @@ -227,7 +227,7 @@ std::mutex& host_mr_mutex() return map_lock; } -CUDF_EXPORT auto& host_mr() +auto& host_mr() { static rmm::host_async_resource_ref host_mr = default_pinned_mr(); return host_mr; From a244d7c81c058048f45270bff52b131329cc91ab Mon Sep 17 00:00:00 2001 From: Vukasin Milovanovic Date: Thu, 9 May 2024 11:35:57 -0700 Subject: [PATCH 17/27] non-indexed stream --- cpp/src/io/utilities/config_utils.cpp | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/cpp/src/io/utilities/config_utils.cpp b/cpp/src/io/utilities/config_utils.cpp index fec5ca88a12..da5c9efbf4f 100644 --- a/cpp/src/io/utilities/config_utils.cpp +++ b/cpp/src/io/utilities/config_utils.cpp @@ -103,7 +103,7 @@ class fixed_pinned_pool_memory_resource { host_pooled_mr* pool_{nullptr}; void* pool_begin_{nullptr}; void* pool_end_{nullptr}; - cuda::stream_ref stream_{cudf::detail::global_cuda_stream_pool().get_stream(0).value()}; + cuda::stream_ref stream_{cudf::detail::global_cuda_stream_pool().get_stream().value()}; public: fixed_pinned_pool_memory_resource(size_t size) From 7076e735aed1415789f6699e63cb06c58863a021 Mon Sep 17 00:00:00 2001 From: Nghia Truong <7416935+ttnghia@users.noreply.github.com> Date: Thu, 9 May 2024 15:02:45 -0600 Subject: [PATCH 18/27] Update cpp/src/io/utilities/config_utils.cpp --- cpp/src/io/utilities/config_utils.cpp | 1 + 1 file changed, 1 insertion(+) diff --git a/cpp/src/io/utilities/config_utils.cpp b/cpp/src/io/utilities/config_utils.cpp index da5c9efbf4f..65ba06e6fed 100644 --- a/cpp/src/io/utilities/config_utils.cpp +++ b/cpp/src/io/utilities/config_utils.cpp @@ -129,6 +129,7 @@ class fixed_pinned_pool_memory_resource { return upstream_mr_.allocate_async(bytes, alignment, stream); } + void do_deallocate_async(void* ptr, std::size_t bytes, std::size_t alignment, From 27d30c80e6c2e4b59a7f7f4bf5c8da2b766fa9d5 Mon Sep 17 00:00:00 2001 From: Vukasin Milovanovic Date: Mon, 13 May 2024 11:13:31 -0700 Subject: [PATCH 19/27] make default_pinned_mr cheap to call multiple times --- cpp/src/io/utilities/config_utils.cpp | 29 ++++++++++++++------------- 1 file changed, 15 insertions(+), 14 deletions(-) diff --git a/cpp/src/io/utilities/config_utils.cpp b/cpp/src/io/utilities/config_utils.cpp index da5c9efbf4f..6d5c8a7727b 100644 --- a/cpp/src/io/utilities/config_utils.cpp +++ b/cpp/src/io/utilities/config_utils.cpp @@ -202,21 +202,22 @@ class fixed_pinned_pool_memory_resource { rmm::host_async_resource_ref default_pinned_mr() { - auto const size = []() -> size_t { - if (auto const env_val = getenv("LIBCUDF_PINNED_POOL_SIZE"); env_val != nullptr) { - return std::atol(env_val); - } - - size_t free{}, total{}; - CUDF_CUDA_TRY(cudaMemGetInfo(&free, &total)); - // 0.5% of the total device memory, capped at 100MB - return std::min((total / 200 + 255) & ~255, size_t{100} * 1024 * 1024); - }(); + static fixed_pinned_pool_memory_resource mr = []() { + auto const size = []() -> size_t { + if (auto const env_val = getenv("LIBCUDF_PINNED_POOL_SIZE"); env_val != nullptr) { + return std::atol(env_val); + } - CUDF_LOG_INFO("Pinned pool size = {}", size); + size_t free{}, total{}; + CUDF_CUDA_TRY(cudaMemGetInfo(&free, &total)); + // 0.5% of the total device memory, capped at 100MB + return std::min((total / 200 + 255) & ~255, size_t{100} * 1024 * 1024); + }(); + CUDF_LOG_INFO("Pinned pool size = {}", size); - // make the pool with max size equal to the initial size - static fixed_pinned_pool_memory_resource mr{size}; + // make the pool with max size equal to the initial size + return fixed_pinned_pool_memory_resource{size}; + }(); return mr; } @@ -227,7 +228,7 @@ std::mutex& host_mr_mutex() return map_lock; } -auto& host_mr() +rmm::host_async_resource_ref host_mr() { static rmm::host_async_resource_ref host_mr = default_pinned_mr(); return host_mr; From 224e68fb3d639042a4cf88a9d1f4c36dcc546fd8 Mon Sep 17 00:00:00 2001 From: Vukasin Milovanovic Date: Mon, 13 May 2024 11:16:35 -0700 Subject: [PATCH 20/27] static_assert fixed_pinned_pool_memory_resource --- cpp/src/io/utilities/config_utils.cpp | 5 +++++ 1 file changed, 5 insertions(+) diff --git a/cpp/src/io/utilities/config_utils.cpp b/cpp/src/io/utilities/config_utils.cpp index 6d5c8a7727b..866c75408ca 100644 --- a/cpp/src/io/utilities/config_utils.cpp +++ b/cpp/src/io/utilities/config_utils.cpp @@ -200,6 +200,11 @@ class fixed_pinned_pool_memory_resource { } }; +static_assert(cuda::mr::resource_with, + ""); + rmm::host_async_resource_ref default_pinned_mr() { static fixed_pinned_pool_memory_resource mr = []() { From b2fd734c13de9fbb4043140122d3d892e15151e3 Mon Sep 17 00:00:00 2001 From: Vukasin Milovanovic Date: Mon, 13 May 2024 15:11:16 -0700 Subject: [PATCH 21/27] fix host_mr --- cpp/src/io/utilities/config_utils.cpp | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/cpp/src/io/utilities/config_utils.cpp b/cpp/src/io/utilities/config_utils.cpp index c4f4fbd7b7d..f345b4868f1 100644 --- a/cpp/src/io/utilities/config_utils.cpp +++ b/cpp/src/io/utilities/config_utils.cpp @@ -234,7 +234,7 @@ std::mutex& host_mr_mutex() return map_lock; } -rmm::host_async_resource_ref host_mr() +rmm::host_async_resource_ref& host_mr() { static rmm::host_async_resource_ref host_mr = default_pinned_mr(); return host_mr; From 0eccf9a6656a0dc14108dc29f0afbb1c303eb447 Mon Sep 17 00:00:00 2001 From: Vukasin Milovanovic Date: Mon, 13 May 2024 15:52:52 -0700 Subject: [PATCH 22/27] add config function --- cpp/include/cudf/io/memory_resource.hpp | 9 +++++++++ cpp/src/io/utilities/config_utils.cpp | 24 +++++++++++++++++++----- 2 files changed, 28 insertions(+), 5 deletions(-) diff --git a/cpp/include/cudf/io/memory_resource.hpp b/cpp/include/cudf/io/memory_resource.hpp index ea79d6a3029..f8ef3bfdf93 100644 --- a/cpp/include/cudf/io/memory_resource.hpp +++ b/cpp/include/cudf/io/memory_resource.hpp @@ -41,4 +41,13 @@ rmm::host_async_resource_ref set_host_memory_resource(rmm::host_async_resource_r */ rmm::host_async_resource_ref get_host_memory_resource(); +/** + * @brief Configure the size of the default host memory resource. + * + * Must be called before any other function in this header. + * + * @param size The size of the default host memory resource + */ +void config_host_memory_resource(size_t size); + } // namespace cudf::io diff --git a/cpp/src/io/utilities/config_utils.cpp b/cpp/src/io/utilities/config_utils.cpp index f345b4868f1..a5388a77b65 100644 --- a/cpp/src/io/utilities/config_utils.cpp +++ b/cpp/src/io/utilities/config_utils.cpp @@ -206,14 +206,16 @@ static_assert(cuda::mr::resource_with, ""); -rmm::host_async_resource_ref default_pinned_mr() +rmm::host_async_resource_ref make_default_pinned_mr(std::optional config_size) { - static fixed_pinned_pool_memory_resource mr = []() { - auto const size = []() -> size_t { + static fixed_pinned_pool_memory_resource mr = [config_size]() { + auto const size = [&config_size]() -> size_t { if (auto const env_val = getenv("LIBCUDF_PINNED_POOL_SIZE"); env_val != nullptr) { return std::atol(env_val); } + if (config_size.has_value()) { return *config_size; } + size_t free{}, total{}; CUDF_CUDA_TRY(cudaMemGetInfo(&free, &total)); // 0.5% of the total device memory, capped at 100MB @@ -228,6 +230,12 @@ rmm::host_async_resource_ref default_pinned_mr() return mr; } +rmm::host_async_resource_ref make_host_mr(std::optional size) +{ + static rmm::host_async_resource_ref mr_ref = make_default_pinned_mr(size); + return mr_ref; +} + std::mutex& host_mr_mutex() { static std::mutex map_lock; @@ -236,8 +244,8 @@ std::mutex& host_mr_mutex() rmm::host_async_resource_ref& host_mr() { - static rmm::host_async_resource_ref host_mr = default_pinned_mr(); - return host_mr; + static rmm::host_async_resource_ref mr_ref = make_host_mr(std::nullopt); + return mr_ref; } } // namespace @@ -256,4 +264,10 @@ rmm::host_async_resource_ref get_host_memory_resource() return host_mr(); } +void config_host_memory_resource(size_t size) +{ + std::scoped_lock lock{host_mr_mutex()}; + make_host_mr(size); +} + } // namespace cudf::io From ecd648110b94236310359a995b97acc41e66757e Mon Sep 17 00:00:00 2001 From: Vukasin Milovanovic Date: Mon, 13 May 2024 16:35:36 -0700 Subject: [PATCH 23/27] align config size; add missing header --- cpp/src/io/utilities/config_utils.cpp | 10 +++++++--- 1 file changed, 7 insertions(+), 3 deletions(-) diff --git a/cpp/src/io/utilities/config_utils.cpp b/cpp/src/io/utilities/config_utils.cpp index a5388a77b65..ad32dcf1229 100644 --- a/cpp/src/io/utilities/config_utils.cpp +++ b/cpp/src/io/utilities/config_utils.cpp @@ -17,6 +17,7 @@ #include "config_utils.hpp" #include +#include #include #include @@ -219,12 +220,15 @@ rmm::host_async_resource_ref make_default_pinned_mr(std::optional config size_t free{}, total{}; CUDF_CUDA_TRY(cudaMemGetInfo(&free, &total)); // 0.5% of the total device memory, capped at 100MB - return std::min((total / 200 + 255) & ~255, size_t{100} * 1024 * 1024); + return std::min(total / 200, size_t{100} * 1024 * 1024); }(); - CUDF_LOG_INFO("Pinned pool size = {}", size); + + // rmm requires the pool size to be a multiple of 256 bytes + auto const aligned_size = (size + 255) & ~255; + CUDF_LOG_INFO("Pinned pool size = {}", aligned_size); // make the pool with max size equal to the initial size - return fixed_pinned_pool_memory_resource{size}; + return fixed_pinned_pool_memory_resource{aligned_size}; }(); return mr; From ecb5f5a346ca2b7c6c5fd25ab2ea6963d62e894c Mon Sep 17 00:00:00 2001 From: Vukasin Milovanovic Date: Tue, 14 May 2024 13:07:32 -0700 Subject: [PATCH 24/27] CUDF_EXPORT --- cpp/src/io/utilities/config_utils.cpp | 12 ++++++------ 1 file changed, 6 insertions(+), 6 deletions(-) diff --git a/cpp/src/io/utilities/config_utils.cpp b/cpp/src/io/utilities/config_utils.cpp index ad32dcf1229..eba65281390 100644 --- a/cpp/src/io/utilities/config_utils.cpp +++ b/cpp/src/io/utilities/config_utils.cpp @@ -207,7 +207,9 @@ static_assert(cuda::mr::resource_with, ""); -rmm::host_async_resource_ref make_default_pinned_mr(std::optional config_size) +} // namespace + +CUDF_EXPORT rmm::host_async_resource_ref make_default_pinned_mr(std::optional config_size) { static fixed_pinned_pool_memory_resource mr = [config_size]() { auto const size = [&config_size]() -> size_t { @@ -234,26 +236,24 @@ rmm::host_async_resource_ref make_default_pinned_mr(std::optional config return mr; } -rmm::host_async_resource_ref make_host_mr(std::optional size) +CUDF_EXPORT rmm::host_async_resource_ref make_host_mr(std::optional size) { static rmm::host_async_resource_ref mr_ref = make_default_pinned_mr(size); return mr_ref; } -std::mutex& host_mr_mutex() +CUDF_EXPORT std::mutex& host_mr_mutex() { static std::mutex map_lock; return map_lock; } -rmm::host_async_resource_ref& host_mr() +CUDF_EXPORT rmm::host_async_resource_ref& host_mr() { static rmm::host_async_resource_ref mr_ref = make_host_mr(std::nullopt); return mr_ref; } -} // namespace - rmm::host_async_resource_ref set_host_memory_resource(rmm::host_async_resource_ref mr) { std::scoped_lock lock{host_mr_mutex()}; From f0d0bf08aebce7dfa6f3daded830fc91dcb9d4b2 Mon Sep 17 00:00:00 2001 From: Vukasin Milovanovic Date: Tue, 14 May 2024 14:45:29 -0700 Subject: [PATCH 25/27] fail config if resource is already created --- cpp/include/cudf/io/memory_resource.hpp | 2 +- cpp/src/io/utilities/config_utils.cpp | 22 ++++++++++++++-------- 2 files changed, 15 insertions(+), 9 deletions(-) diff --git a/cpp/include/cudf/io/memory_resource.hpp b/cpp/include/cudf/io/memory_resource.hpp index f8ef3bfdf93..931a83424cb 100644 --- a/cpp/include/cudf/io/memory_resource.hpp +++ b/cpp/include/cudf/io/memory_resource.hpp @@ -48,6 +48,6 @@ rmm::host_async_resource_ref get_host_memory_resource(); * * @param size The size of the default host memory resource */ -void config_host_memory_resource(size_t size); +void config_default_host_memory_resource(size_t size); } // namespace cudf::io diff --git a/cpp/src/io/utilities/config_utils.cpp b/cpp/src/io/utilities/config_utils.cpp index eba65281390..a70a987ea8f 100644 --- a/cpp/src/io/utilities/config_utils.cpp +++ b/cpp/src/io/utilities/config_utils.cpp @@ -209,7 +209,7 @@ static_assert(cuda::mr::resource_with config_size) +CUDF_EXPORT rmm::host_async_resource_ref& make_default_pinned_mr(std::optional config_size) { static fixed_pinned_pool_memory_resource mr = [config_size]() { auto const size = [&config_size]() -> size_t { @@ -233,12 +233,7 @@ CUDF_EXPORT rmm::host_async_resource_ref make_default_pinned_mr(std::optional size) -{ - static rmm::host_async_resource_ref mr_ref = make_default_pinned_mr(size); + static rmm::host_async_resource_ref mr_ref{mr}; return mr_ref; } @@ -248,6 +243,17 @@ CUDF_EXPORT std::mutex& host_mr_mutex() return map_lock; } +// Must be called with the host_mr_mutex mutex held +CUDF_EXPORT rmm::host_async_resource_ref& make_host_mr(std::optional size) +{ + static rmm::host_async_resource_ref* mr_ref = nullptr; + CUDF_EXPECTS(mr_ref == nullptr, "The default host memory resource has already been created"); + mr_ref = &make_default_pinned_mr(size); + + return *mr_ref; +} + +// Must be called with the host_mr_mutex mutex held CUDF_EXPORT rmm::host_async_resource_ref& host_mr() { static rmm::host_async_resource_ref mr_ref = make_host_mr(std::nullopt); @@ -268,7 +274,7 @@ rmm::host_async_resource_ref get_host_memory_resource() return host_mr(); } -void config_host_memory_resource(size_t size) +void config_default_host_memory_resource(size_t size) { std::scoped_lock lock{host_mr_mutex()}; make_host_mr(size); From f989a56b72f456c24dadbd18c116ba2abd4a1fc9 Mon Sep 17 00:00:00 2001 From: Vukasin Milovanovic Date: Wed, 15 May 2024 14:34:43 -0700 Subject: [PATCH 26/27] fix config check --- cpp/include/cudf/io/memory_resource.hpp | 13 ++++++++++--- cpp/src/io/utilities/config_utils.cpp | 14 +++++++++----- 2 files changed, 19 insertions(+), 8 deletions(-) diff --git a/cpp/include/cudf/io/memory_resource.hpp b/cpp/include/cudf/io/memory_resource.hpp index 931a83424cb..9518dbcae78 100644 --- a/cpp/include/cudf/io/memory_resource.hpp +++ b/cpp/include/cudf/io/memory_resource.hpp @@ -18,6 +18,8 @@ #include +#include + namespace cudf::io { /** @@ -41,13 +43,18 @@ rmm::host_async_resource_ref set_host_memory_resource(rmm::host_async_resource_r */ rmm::host_async_resource_ref get_host_memory_resource(); +// Options to configure the default host memory resource +struct host_mr_options { + std::optional pool_size; +}; + /** * @brief Configure the size of the default host memory resource. * - * Must be called before any other function in this header. + * @throws cudf::logic_error if called after the default host memory resource has been created * - * @param size The size of the default host memory resource + * @param opts Options to configure the default host memory resource */ -void config_default_host_memory_resource(size_t size); +void config_default_host_memory_resource(host_mr_options const& opts); } // namespace cudf::io diff --git a/cpp/src/io/utilities/config_utils.cpp b/cpp/src/io/utilities/config_utils.cpp index a70a987ea8f..7720c073a97 100644 --- a/cpp/src/io/utilities/config_utils.cpp +++ b/cpp/src/io/utilities/config_utils.cpp @@ -244,11 +244,15 @@ CUDF_EXPORT std::mutex& host_mr_mutex() } // Must be called with the host_mr_mutex mutex held -CUDF_EXPORT rmm::host_async_resource_ref& make_host_mr(std::optional size) +CUDF_EXPORT rmm::host_async_resource_ref& make_host_mr(std::optional const& opts) { static rmm::host_async_resource_ref* mr_ref = nullptr; - CUDF_EXPECTS(mr_ref == nullptr, "The default host memory resource has already been created"); - mr_ref = &make_default_pinned_mr(size); + if (mr_ref == nullptr) { + mr_ref = &make_default_pinned_mr(opts ? opts->pool_size : std::nullopt); + } else { + // Throw an error if the user tries to reconfigure the default host resource + CUDF_EXPECTS(opts == std::nullopt, "The default host memory resource has already been created"); + } return *mr_ref; } @@ -274,10 +278,10 @@ rmm::host_async_resource_ref get_host_memory_resource() return host_mr(); } -void config_default_host_memory_resource(size_t size) +void config_default_host_memory_resource(host_mr_options const& opts) { std::scoped_lock lock{host_mr_mutex()}; - make_host_mr(size); + make_host_mr(opts); } } // namespace cudf::io From 2b4952a5732743b83f3332e9a54788a6d79334f8 Mon Sep 17 00:00:00 2001 From: Vukasin Milovanovic Date: Wed, 15 May 2024 15:02:32 -0700 Subject: [PATCH 27/27] docs --- cpp/include/cudf/io/memory_resource.hpp | 7 +++++-- 1 file changed, 5 insertions(+), 2 deletions(-) diff --git a/cpp/include/cudf/io/memory_resource.hpp b/cpp/include/cudf/io/memory_resource.hpp index 9518dbcae78..e31ebce4b1f 100644 --- a/cpp/include/cudf/io/memory_resource.hpp +++ b/cpp/include/cudf/io/memory_resource.hpp @@ -43,9 +43,12 @@ rmm::host_async_resource_ref set_host_memory_resource(rmm::host_async_resource_r */ rmm::host_async_resource_ref get_host_memory_resource(); -// Options to configure the default host memory resource +/** + * @brief Options to configure the default host memory resource + */ struct host_mr_options { - std::optional pool_size; + std::optional pool_size; ///< The size of the pool to use for the default host memory + ///< resource. If not set, the default pool size is used. }; /**