From 4eaaee6cc520bab866e75cac6bdc480e26699d19 Mon Sep 17 00:00:00 2001 From: Alessandro Bellina Date: Wed, 22 May 2024 08:20:52 -0700 Subject: [PATCH 01/10] Pass default host pool parameters as part of set and get --- cpp/include/cudf/io/memory_resource.hpp | 41 +++++++++++++------------ cpp/src/io/utilities/config_utils.cpp | 38 +++++++---------------- 2 files changed, 33 insertions(+), 46 deletions(-) diff --git a/cpp/include/cudf/io/memory_resource.hpp b/cpp/include/cudf/io/memory_resource.hpp index e31ebce4b1f..0ee86b015a9 100644 --- a/cpp/include/cudf/io/memory_resource.hpp +++ b/cpp/include/cudf/io/memory_resource.hpp @@ -22,6 +22,14 @@ namespace cudf::io { +/** + * @brief Options to configure the default host memory resource + */ +struct host_mr_options { + 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. +}; + /** * @brief Set the rmm resource to be used for host memory allocations by * cudf::detail::hostdevice_vector @@ -30,34 +38,29 @@ namespace cudf::io { * bouncing state between the cpu and the gpu. The resource set with this function (typically a * pinned memory allocator) is what it uses to allocate space for it's host-side buffer. * + * The default_opts parameter allows the caller to customize the default host memory resource + * if it hasn't been configured already, otherwise the argument is ignored. + * Omitting this argument (nullopt) means cuDF will use defaults to initialize a host pinned pool. + * * @param mr The rmm resource to be used for host-side allocations + * @param default_opts Options to configure the default host memory resource * @return The previous resource that was in use */ -rmm::host_async_resource_ref set_host_memory_resource(rmm::host_async_resource_ref mr); +rmm::host_async_resource_ref set_host_memory_resource( + rmm::host_async_resource_ref mr, std::optional const& default_opts = nullopt); /** * @brief Get the rmm resource being used for host memory allocations by * cudf::detail::hostdevice_vector * - * @return The rmm resource used for host-side allocations - */ -rmm::host_async_resource_ref get_host_memory_resource(); - -/** - * @brief Options to configure the default host memory resource - */ -struct host_mr_options { - 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. -}; - -/** - * @brief Configure the size of the default host memory resource. + * The default_opts parameter allows the caller to customize the default host memory resource + * if it hasn't been configured already, otherwise the argument is ignored. + * Omitting this argument (nullopt) means cuDF will use defaults to initialize a host pinned pool. * - * @throws cudf::logic_error if called after the default host memory resource has been created - * - * @param opts Options to configure the default host memory resource + * @param default_opts Options to configure the default host memory resource + * @return The rmm resource used for host-side allocations */ -void config_default_host_memory_resource(host_mr_options const& opts); +rmm::host_async_resource_ref get_host_memory_resource( + std::optional const& default_opts = nullopt); } // namespace cudf::io diff --git a/cpp/src/io/utilities/config_utils.cpp b/cpp/src/io/utilities/config_utils.cpp index 7720c073a97..0d58066d328 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 { @@ -244,44 +244,28 @@ 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 const& opts) +CUDF_EXPORT rmm::host_async_resource_ref& host_mr( + std::optional const& default_opts) { - static rmm::host_async_resource_ref* mr_ref = nullptr; - 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; -} - -// 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); + static rmm::host_async_resource_ref mr_ref = + make_default_pinned_mr(default_opts ? default_opts->pool_size : std::nullopt); return mr_ref; } -rmm::host_async_resource_ref set_host_memory_resource(rmm::host_async_resource_ref mr) +rmm::host_async_resource_ref set_host_memory_resource( + rmm::host_async_resource_ref mr, std::optional const& default_opts) { std::scoped_lock lock{host_mr_mutex()}; - auto last_mr = host_mr(); + auto last_mr = host_mr(default_opts); host_mr() = mr; return last_mr; } -rmm::host_async_resource_ref get_host_memory_resource() -{ - std::scoped_lock lock{host_mr_mutex()}; - return host_mr(); -} - -void config_default_host_memory_resource(host_mr_options const& opts) +rmm::host_async_resource_ref get_host_memory_resource( + std::optional const& default_opts) { std::scoped_lock lock{host_mr_mutex()}; - make_host_mr(opts); + return host_mr(default_opts); } } // namespace cudf::io From ff392b56c8c4c99ad3e96d657a9fe998c8a4b794 Mon Sep 17 00:00:00 2001 From: Alessandro Bellina Date: Wed, 22 May 2024 08:22:37 -0700 Subject: [PATCH 02/10] Revert "[JNI] Expose java API for cudf::io::config_host_memory_resource (#15745)" This reverts commit 8b7245548c63d1ce84031a0bd187cbfb8e072f8c. --- .../main/java/ai/rapids/cudf/PinnedMemoryPool.java | 13 ------------- java/src/main/java/ai/rapids/cudf/Rmm.java | 10 ---------- java/src/main/native/src/RmmJni.cpp | 11 ----------- 3 files changed, 34 deletions(-) diff --git a/java/src/main/java/ai/rapids/cudf/PinnedMemoryPool.java b/java/src/main/java/ai/rapids/cudf/PinnedMemoryPool.java index 9038700cb30..6cb34683e5a 100644 --- a/java/src/main/java/ai/rapids/cudf/PinnedMemoryPool.java +++ b/java/src/main/java/ai/rapids/cudf/PinnedMemoryPool.java @@ -252,17 +252,4 @@ private synchronized HostMemoryBuffer tryAllocateInternal(long bytes) { private synchronized void free(long address, long size) { Rmm.freeFromPinnedPool(this.poolHandle, address, size); } - - /** - * Sets the size of the cuDF default pinned pool. - * - * @note This has to be called before cuDF functions are executed. - * - * @param size initial and maximum size for the cuDF default pinned pool. - * Pass size=0 to disable the default pool. - */ - public static synchronized void configureDefaultCudfPinnedPoolSize(long size) { - Rmm.configureDefaultCudfPinnedPoolSize(size); - } - } diff --git a/java/src/main/java/ai/rapids/cudf/Rmm.java b/java/src/main/java/ai/rapids/cudf/Rmm.java index fdbdfdfff6f..6e9f90e477f 100755 --- a/java/src/main/java/ai/rapids/cudf/Rmm.java +++ b/java/src/main/java/ai/rapids/cudf/Rmm.java @@ -266,16 +266,6 @@ public static synchronized void initialize(int allocationMode, LogConf logConf, } } - /** - * Sets the size of the cuDF default pinned pool. - * - * @note This has to be called before cuDF functions are executed. - * - * @param size initial and maximum size for the cuDF default pinned pool. - * Pass size=0 to disable the default pool. - */ - public static synchronized native void configureDefaultCudfPinnedPoolSize(long size); - /** * Get the most recently set pool size or -1 if RMM has not been initialized or pooling is * not enabled. diff --git a/java/src/main/native/src/RmmJni.cpp b/java/src/main/native/src/RmmJni.cpp index 9c015fee409..68453c924d6 100644 --- a/java/src/main/native/src/RmmJni.cpp +++ b/java/src/main/native/src/RmmJni.cpp @@ -1106,15 +1106,4 @@ JNIEXPORT void JNICALL Java_ai_rapids_cudf_Rmm_freeFromFallbackPinnedPool(JNIEnv } CATCH_STD(env, ) } - -JNIEXPORT void JNICALL Java_ai_rapids_cudf_Rmm_configureDefaultCudfPinnedPoolSize(JNIEnv* env, - jclass clazz, - jlong size) -{ - try { - cudf::jni::auto_set_device(env); - cudf::io::config_default_host_memory_resource(cudf::io::host_mr_options{size}); - } - CATCH_STD(env, ) -} } From 068023910f674f70e6b2581083ce19f1a8d2dfc4 Mon Sep 17 00:00:00 2001 From: Alessandro Bellina Date: Wed, 22 May 2024 08:30:46 -0700 Subject: [PATCH 03/10] nullopt -> std::nullopt --- cpp/include/cudf/io/memory_resource.hpp | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/cpp/include/cudf/io/memory_resource.hpp b/cpp/include/cudf/io/memory_resource.hpp index 0ee86b015a9..b26c843a45f 100644 --- a/cpp/include/cudf/io/memory_resource.hpp +++ b/cpp/include/cudf/io/memory_resource.hpp @@ -47,7 +47,7 @@ struct host_mr_options { * @return The previous resource that was in use */ rmm::host_async_resource_ref set_host_memory_resource( - rmm::host_async_resource_ref mr, std::optional const& default_opts = nullopt); + rmm::host_async_resource_ref mr, std::optional const& default_opts = std::nullopt); /** * @brief Get the rmm resource being used for host memory allocations by @@ -61,6 +61,6 @@ rmm::host_async_resource_ref set_host_memory_resource( * @return The rmm resource used for host-side allocations */ rmm::host_async_resource_ref get_host_memory_resource( - std::optional const& default_opts = nullopt); + std::optional const& default_opts = std::nullopt); } // namespace cudf::io From 5b3e7a268b0f629e5123eecd19bb93db01491a96 Mon Sep 17 00:00:00 2001 From: Alessandro Bellina Date: Wed, 22 May 2024 08:37:37 -0700 Subject: [PATCH 04/10] Pass std::nullopt to host_mr when setting its reference --- 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 0d58066d328..b14d9f3e88e 100644 --- a/cpp/src/io/utilities/config_utils.cpp +++ b/cpp/src/io/utilities/config_utils.cpp @@ -257,7 +257,7 @@ rmm::host_async_resource_ref set_host_memory_resource( { std::scoped_lock lock{host_mr_mutex()}; auto last_mr = host_mr(default_opts); - host_mr() = mr; + host_mr(std::nullopt) = mr; return last_mr; } From fd0d2b4ebec458f44808ebea804f3146039373f8 Mon Sep 17 00:00:00 2001 From: Alessandro Bellina Date: Wed, 22 May 2024 08:41:16 -0700 Subject: [PATCH 05/10] Use host_mr_options in JNI --- cpp/include/cudf/io/memory_resource.hpp | 3 ++- cpp/src/io/utilities/config_utils.cpp | 4 ++-- java/src/main/native/src/RmmJni.cpp | 4 +++- 3 files changed, 7 insertions(+), 4 deletions(-) diff --git a/cpp/include/cudf/io/memory_resource.hpp b/cpp/include/cudf/io/memory_resource.hpp index b26c843a45f..070af4f8588 100644 --- a/cpp/include/cudf/io/memory_resource.hpp +++ b/cpp/include/cudf/io/memory_resource.hpp @@ -47,7 +47,8 @@ struct host_mr_options { * @return The previous resource that was in use */ rmm::host_async_resource_ref set_host_memory_resource( - rmm::host_async_resource_ref mr, std::optional const& default_opts = std::nullopt); + rmm::host_async_resource_ref mr, + std::optional const& default_opts = std::nullopt); /** * @brief Get the rmm resource being used for host memory allocations by diff --git a/cpp/src/io/utilities/config_utils.cpp b/cpp/src/io/utilities/config_utils.cpp index b14d9f3e88e..93a4c8b9a71 100644 --- a/cpp/src/io/utilities/config_utils.cpp +++ b/cpp/src/io/utilities/config_utils.cpp @@ -256,8 +256,8 @@ rmm::host_async_resource_ref set_host_memory_resource( rmm::host_async_resource_ref mr, std::optional const& default_opts) { std::scoped_lock lock{host_mr_mutex()}; - auto last_mr = host_mr(default_opts); - host_mr(std::nullopt) = mr; + auto last_mr = host_mr(default_opts); + host_mr(std::nullopt) = mr; return last_mr; } diff --git a/java/src/main/native/src/RmmJni.cpp b/java/src/main/native/src/RmmJni.cpp index 68453c924d6..bcc58cff982 100644 --- a/java/src/main/native/src/RmmJni.cpp +++ b/java/src/main/native/src/RmmJni.cpp @@ -1036,7 +1036,9 @@ JNIEXPORT void JNICALL Java_ai_rapids_cudf_Rmm_setCuioPinnedPoolMemoryResource(J // if the regular pinned pool is exhausted pinned_fallback_mr.reset(new pinned_fallback_host_memory_resource(pool)); // set the cuio host mr and store the prior resource in our static variable - prior_cuio_host_mr() = cudf::io::set_host_memory_resource(*pinned_fallback_mr); + // we pass host_mr_options{0} so we disable any default pinned pool creation from cuDF + prior_cuio_host_mr() = + cudf::io::set_host_memory_resource(*pinned_fallback_mr, cudf::io::host_mr_options{0}); } CATCH_STD(env, ) } From 8365382c2215da8b7c4fc851f3821f268786fb96 Mon Sep 17 00:00:00 2001 From: Alessandro Bellina Date: Wed, 22 May 2024 11:32:34 -0700 Subject: [PATCH 06/10] Revert "Revert "[JNI] Expose java API for cudf::io::config_host_memory_resource (#15745)"" This reverts commit ff392b56c8c4c99ad3e96d657a9fe998c8a4b794. --- .../main/java/ai/rapids/cudf/PinnedMemoryPool.java | 13 +++++++++++++ java/src/main/java/ai/rapids/cudf/Rmm.java | 10 ++++++++++ java/src/main/native/src/RmmJni.cpp | 11 +++++++++++ 3 files changed, 34 insertions(+) diff --git a/java/src/main/java/ai/rapids/cudf/PinnedMemoryPool.java b/java/src/main/java/ai/rapids/cudf/PinnedMemoryPool.java index 6cb34683e5a..9038700cb30 100644 --- a/java/src/main/java/ai/rapids/cudf/PinnedMemoryPool.java +++ b/java/src/main/java/ai/rapids/cudf/PinnedMemoryPool.java @@ -252,4 +252,17 @@ private synchronized HostMemoryBuffer tryAllocateInternal(long bytes) { private synchronized void free(long address, long size) { Rmm.freeFromPinnedPool(this.poolHandle, address, size); } + + /** + * Sets the size of the cuDF default pinned pool. + * + * @note This has to be called before cuDF functions are executed. + * + * @param size initial and maximum size for the cuDF default pinned pool. + * Pass size=0 to disable the default pool. + */ + public static synchronized void configureDefaultCudfPinnedPoolSize(long size) { + Rmm.configureDefaultCudfPinnedPoolSize(size); + } + } diff --git a/java/src/main/java/ai/rapids/cudf/Rmm.java b/java/src/main/java/ai/rapids/cudf/Rmm.java index 6e9f90e477f..fdbdfdfff6f 100755 --- a/java/src/main/java/ai/rapids/cudf/Rmm.java +++ b/java/src/main/java/ai/rapids/cudf/Rmm.java @@ -266,6 +266,16 @@ public static synchronized void initialize(int allocationMode, LogConf logConf, } } + /** + * Sets the size of the cuDF default pinned pool. + * + * @note This has to be called before cuDF functions are executed. + * + * @param size initial and maximum size for the cuDF default pinned pool. + * Pass size=0 to disable the default pool. + */ + public static synchronized native void configureDefaultCudfPinnedPoolSize(long size); + /** * Get the most recently set pool size or -1 if RMM has not been initialized or pooling is * not enabled. diff --git a/java/src/main/native/src/RmmJni.cpp b/java/src/main/native/src/RmmJni.cpp index bcc58cff982..18a05b8d0b7 100644 --- a/java/src/main/native/src/RmmJni.cpp +++ b/java/src/main/native/src/RmmJni.cpp @@ -1108,4 +1108,15 @@ JNIEXPORT void JNICALL Java_ai_rapids_cudf_Rmm_freeFromFallbackPinnedPool(JNIEnv } CATCH_STD(env, ) } + +JNIEXPORT void JNICALL Java_ai_rapids_cudf_Rmm_configureDefaultCudfPinnedPoolSize(JNIEnv* env, + jclass clazz, + jlong size) +{ + try { + cudf::jni::auto_set_device(env); + cudf::io::config_default_host_memory_resource(cudf::io::host_mr_options{size}); + } + CATCH_STD(env, ) +} } From 369aa8935bf05a16b7384c4854a6e47e62937cf4 Mon Sep 17 00:00:00 2001 From: Alessandro Bellina Date: Wed, 22 May 2024 12:06:06 -0700 Subject: [PATCH 07/10] Bring back config_default_host_memory_resource but change semantics so it doesnt throw --- cpp/include/cudf/io/memory_resource.hpp | 42 ++++++++--------- cpp/src/io/utilities/config_utils.cpp | 47 ++++++++++++++----- .../java/ai/rapids/cudf/PinnedMemoryPool.java | 7 ++- java/src/main/java/ai/rapids/cudf/Rmm.java | 5 +- java/src/main/native/src/RmmJni.cpp | 15 +++--- 5 files changed, 69 insertions(+), 47 deletions(-) diff --git a/cpp/include/cudf/io/memory_resource.hpp b/cpp/include/cudf/io/memory_resource.hpp index 070af4f8588..d7556df03a4 100644 --- a/cpp/include/cudf/io/memory_resource.hpp +++ b/cpp/include/cudf/io/memory_resource.hpp @@ -22,14 +22,6 @@ namespace cudf::io { -/** - * @brief Options to configure the default host memory resource - */ -struct host_mr_options { - 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. -}; - /** * @brief Set the rmm resource to be used for host memory allocations by * cudf::detail::hostdevice_vector @@ -38,30 +30,34 @@ struct host_mr_options { * bouncing state between the cpu and the gpu. The resource set with this function (typically a * pinned memory allocator) is what it uses to allocate space for it's host-side buffer. * - * The default_opts parameter allows the caller to customize the default host memory resource - * if it hasn't been configured already, otherwise the argument is ignored. - * Omitting this argument (nullopt) means cuDF will use defaults to initialize a host pinned pool. - * * @param mr The rmm resource to be used for host-side allocations - * @param default_opts Options to configure the default host memory resource * @return The previous resource that was in use */ -rmm::host_async_resource_ref set_host_memory_resource( - rmm::host_async_resource_ref mr, - std::optional const& default_opts = std::nullopt); +rmm::host_async_resource_ref set_host_memory_resource(rmm::host_async_resource_ref mr); /** * @brief Get the rmm resource being used for host memory allocations by * cudf::detail::hostdevice_vector * - * The default_opts parameter allows the caller to customize the default host memory resource - * if it hasn't been configured already, otherwise the argument is ignored. - * Omitting this argument (nullopt) means cuDF will use defaults to initialize a host pinned pool. - * - * @param default_opts Options to configure the default host memory resource * @return The rmm resource used for host-side allocations */ -rmm::host_async_resource_ref get_host_memory_resource( - std::optional const& default_opts = std::nullopt); +rmm::host_async_resource_ref get_host_memory_resource(); + +/** + * @brief Options to configure the default host memory resource + */ +struct host_mr_options { + 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. +}; + +/** + * @brief Configure the size of the default host memory resource. + * + * @throws cudf::logic_error if called after the default host memory resource has been created + * + * @param opts Options to configure the default host memory resource + */ +bool 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 93a4c8b9a71..012d308de4c 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 { @@ -244,28 +244,51 @@ CUDF_EXPORT std::mutex& host_mr_mutex() } // Must be called with the host_mr_mutex mutex held -CUDF_EXPORT rmm::host_async_resource_ref& host_mr( - std::optional const& default_opts) +CUDF_EXPORT rmm::host_async_resource_ref& make_host_mr(std::optional const& opts, bool* did_configure = nullptr) { - static rmm::host_async_resource_ref mr_ref = - make_default_pinned_mr(default_opts ? default_opts->pool_size : std::nullopt); + static rmm::host_async_resource_ref* mr_ref = nullptr; + bool configured = false; + if (mr_ref == nullptr) { + configured = true; + mr_ref = &make_default_pinned_mr(opts ? opts->pool_size : std::nullopt); + } + + // 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; +} + +// 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); return mr_ref; } -rmm::host_async_resource_ref set_host_memory_resource( - rmm::host_async_resource_ref mr, std::optional const& default_opts) +rmm::host_async_resource_ref set_host_memory_resource(rmm::host_async_resource_ref mr) { std::scoped_lock lock{host_mr_mutex()}; - auto last_mr = host_mr(default_opts); - host_mr(std::nullopt) = mr; + auto last_mr = host_mr(); + host_mr() = mr; return last_mr; } -rmm::host_async_resource_ref get_host_memory_resource( - std::optional const& default_opts) +rmm::host_async_resource_ref get_host_memory_resource() +{ + std::scoped_lock lock{host_mr_mutex()}; + return host_mr(); +} + +bool config_default_host_memory_resource(host_mr_options const& opts) { std::scoped_lock lock{host_mr_mutex()}; - return host_mr(default_opts); + auto did_configure = false; + make_host_mr(opts, &did_configure); + return did_configure; } } // namespace cudf::io diff --git a/java/src/main/java/ai/rapids/cudf/PinnedMemoryPool.java b/java/src/main/java/ai/rapids/cudf/PinnedMemoryPool.java index 9038700cb30..0157928c522 100644 --- a/java/src/main/java/ai/rapids/cudf/PinnedMemoryPool.java +++ b/java/src/main/java/ai/rapids/cudf/PinnedMemoryPool.java @@ -260,9 +260,12 @@ private synchronized void free(long address, long size) { * * @param size initial and maximum size for the cuDF default pinned pool. * Pass size=0 to disable the default pool. + * + * @return true if we were able to setup the default resource, false if there was + * a resource already set. */ - public static synchronized void configureDefaultCudfPinnedPoolSize(long size) { - Rmm.configureDefaultCudfPinnedPoolSize(size); + public static synchronized boolean configureDefaultCudfPinnedPoolSize(long size) { + return Rmm.configureDefaultCudfPinnedPoolSize(size); } } diff --git a/java/src/main/java/ai/rapids/cudf/Rmm.java b/java/src/main/java/ai/rapids/cudf/Rmm.java index fdbdfdfff6f..a60c19b0baf 100755 --- a/java/src/main/java/ai/rapids/cudf/Rmm.java +++ b/java/src/main/java/ai/rapids/cudf/Rmm.java @@ -273,8 +273,11 @@ public static synchronized void initialize(int allocationMode, LogConf logConf, * * @param size initial and maximum size for the cuDF default pinned pool. * Pass size=0 to disable the default pool. + * + * @return true if we were able to setup the default resource, false if there was + * a resource already set. */ - public static synchronized native void configureDefaultCudfPinnedPoolSize(long size); + public static synchronized native boolean configureDefaultCudfPinnedPoolSize(long size); /** * Get the most recently set pool size or -1 if RMM has not been initialized or pooling is diff --git a/java/src/main/native/src/RmmJni.cpp b/java/src/main/native/src/RmmJni.cpp index 18a05b8d0b7..fa78f6ca4e2 100644 --- a/java/src/main/native/src/RmmJni.cpp +++ b/java/src/main/native/src/RmmJni.cpp @@ -1035,10 +1035,7 @@ JNIEXPORT void JNICALL Java_ai_rapids_cudf_Rmm_setCuioPinnedPoolMemoryResource(J // create a pinned fallback pool that will allocate pinned memory // if the regular pinned pool is exhausted pinned_fallback_mr.reset(new pinned_fallback_host_memory_resource(pool)); - // set the cuio host mr and store the prior resource in our static variable - // we pass host_mr_options{0} so we disable any default pinned pool creation from cuDF - prior_cuio_host_mr() = - cudf::io::set_host_memory_resource(*pinned_fallback_mr, cudf::io::host_mr_options{0}); + prior_cuio_host_mr() = cudf::io::set_host_memory_resource(*pinned_fallback_mr); } CATCH_STD(env, ) } @@ -1109,14 +1106,14 @@ JNIEXPORT void JNICALL Java_ai_rapids_cudf_Rmm_freeFromFallbackPinnedPool(JNIEnv CATCH_STD(env, ) } -JNIEXPORT void JNICALL Java_ai_rapids_cudf_Rmm_configureDefaultCudfPinnedPoolSize(JNIEnv* env, - jclass clazz, - jlong size) +JNIEXPORT jboolean JNICALL Java_ai_rapids_cudf_Rmm_configureDefaultCudfPinnedPoolSize(JNIEnv* env, + jclass clazz, + jlong size) { try { cudf::jni::auto_set_device(env); - cudf::io::config_default_host_memory_resource(cudf::io::host_mr_options{size}); + return cudf::io::config_default_host_memory_resource(cudf::io::host_mr_options{size}); } - CATCH_STD(env, ) + CATCH_STD(env, false) } } From 7ac43022a346cbe0f50b1644da55a666152a5159 Mon Sep 17 00:00:00 2001 From: Alessandro Bellina Date: Wed, 22 May 2024 12:11:38 -0700 Subject: [PATCH 08/10] Doxygen/code style --- cpp/include/cudf/io/memory_resource.hpp | 2 ++ cpp/src/io/utilities/config_utils.cpp | 11 +++++------ .../main/java/ai/rapids/cudf/PinnedMemoryPool.java | 2 +- java/src/main/java/ai/rapids/cudf/Rmm.java | 2 +- 4 files changed, 9 insertions(+), 8 deletions(-) diff --git a/cpp/include/cudf/io/memory_resource.hpp b/cpp/include/cudf/io/memory_resource.hpp index d7556df03a4..a36e220ae7b 100644 --- a/cpp/include/cudf/io/memory_resource.hpp +++ b/cpp/include/cudf/io/memory_resource.hpp @@ -57,6 +57,8 @@ struct host_mr_options { * @throws cudf::logic_error if called after the default host memory resource has been created * * @param opts Options to configure the default host memory resource + * @return True if this call successfully configured the host memory resource, false if a + * a resource was already configured. */ bool config_default_host_memory_resource(host_mr_options const& opts); diff --git a/cpp/src/io/utilities/config_utils.cpp b/cpp/src/io/utilities/config_utils.cpp index 012d308de4c..dad1135e766 100644 --- a/cpp/src/io/utilities/config_utils.cpp +++ b/cpp/src/io/utilities/config_utils.cpp @@ -244,20 +244,19 @@ 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 const& opts, bool* did_configure = nullptr) +CUDF_EXPORT rmm::host_async_resource_ref& make_host_mr(std::optional const& opts, + bool* did_configure = nullptr) { static rmm::host_async_resource_ref* mr_ref = nullptr; - bool configured = false; + bool configured = false; if (mr_ref == nullptr) { configured = true; - mr_ref = &make_default_pinned_mr(opts ? opts->pool_size : std::nullopt); + mr_ref = &make_default_pinned_mr(opts ? opts->pool_size : std::nullopt); } // 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; - } + if (did_configure != nullptr) { *did_configure = configured; } return *mr_ref; } diff --git a/java/src/main/java/ai/rapids/cudf/PinnedMemoryPool.java b/java/src/main/java/ai/rapids/cudf/PinnedMemoryPool.java index 0157928c522..83b801db7fb 100644 --- a/java/src/main/java/ai/rapids/cudf/PinnedMemoryPool.java +++ b/java/src/main/java/ai/rapids/cudf/PinnedMemoryPool.java @@ -260,7 +260,7 @@ private synchronized void free(long address, long size) { * * @param size initial and maximum size for the cuDF default pinned pool. * Pass size=0 to disable the default pool. - * + * * @return true if we were able to setup the default resource, false if there was * a resource already set. */ diff --git a/java/src/main/java/ai/rapids/cudf/Rmm.java b/java/src/main/java/ai/rapids/cudf/Rmm.java index a60c19b0baf..4dee1b7aa24 100755 --- a/java/src/main/java/ai/rapids/cudf/Rmm.java +++ b/java/src/main/java/ai/rapids/cudf/Rmm.java @@ -273,7 +273,7 @@ public static synchronized void initialize(int allocationMode, LogConf logConf, * * @param size initial and maximum size for the cuDF default pinned pool. * Pass size=0 to disable the default pool. - * + * * @return true if we were able to setup the default resource, false if there was * a resource already set. */ From afd4cbe6f815faee36fcb977e0b86b866a257351 Mon Sep 17 00:00:00 2001 From: Alessandro Bellina Date: Wed, 22 May 2024 12:44:23 -0700 Subject: [PATCH 09/10] Return std::pair instead of an out pointer --- cpp/src/io/utilities/config_utils.cpp | 18 ++++++------------ 1 file changed, 6 insertions(+), 12 deletions(-) diff --git a/cpp/src/io/utilities/config_utils.cpp b/cpp/src/io/utilities/config_utils.cpp index dad1135e766..f7331763ede 100644 --- a/cpp/src/io/utilities/config_utils.cpp +++ b/cpp/src/io/utilities/config_utils.cpp @@ -244,8 +244,8 @@ 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 const& opts, - bool* did_configure = nullptr) +CUDF_EXPORT std::pair make_host_mr( + std::optional const& opts) { static rmm::host_async_resource_ref* mr_ref = nullptr; bool configured = false; @@ -253,18 +253,13 @@ CUDF_EXPORT rmm::host_async_resource_ref& make_host_mr(std::optionalpool_size : std::nullopt); } - - // 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; + return std::make_pair(mr_ref, configured); } // 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); + static rmm::host_async_resource_ref mr_ref = *std::get<0>(make_host_mr(std::nullopt)); return mr_ref; } @@ -285,9 +280,8 @@ rmm::host_async_resource_ref get_host_memory_resource() bool config_default_host_memory_resource(host_mr_options const& opts) { std::scoped_lock lock{host_mr_mutex()}; - auto did_configure = false; - make_host_mr(opts, &did_configure); - return did_configure; + auto configured = std::get<1>(make_host_mr(opts)); + return configured; } } // namespace cudf::io From d1ccbafe4aec4cbe236835646bef2df427286ee2 Mon Sep 17 00:00:00 2001 From: Alessandro Bellina Date: Wed, 22 May 2024 14:21:50 -0700 Subject: [PATCH 10/10] Revert "Return std::pair instead of an out pointer" This reverts commit afd4cbe6f815faee36fcb977e0b86b866a257351. --- cpp/src/io/utilities/config_utils.cpp | 18 ++++++++++++------ 1 file changed, 12 insertions(+), 6 deletions(-) diff --git a/cpp/src/io/utilities/config_utils.cpp b/cpp/src/io/utilities/config_utils.cpp index f7331763ede..dad1135e766 100644 --- a/cpp/src/io/utilities/config_utils.cpp +++ b/cpp/src/io/utilities/config_utils.cpp @@ -244,8 +244,8 @@ CUDF_EXPORT std::mutex& host_mr_mutex() } // Must be called with the host_mr_mutex mutex held -CUDF_EXPORT std::pair make_host_mr( - std::optional const& opts) +CUDF_EXPORT rmm::host_async_resource_ref& make_host_mr(std::optional const& opts, + bool* did_configure = nullptr) { static rmm::host_async_resource_ref* mr_ref = nullptr; bool configured = false; @@ -253,13 +253,18 @@ CUDF_EXPORT std::pair make_host_mr( configured = true; mr_ref = &make_default_pinned_mr(opts ? opts->pool_size : std::nullopt); } - return std::make_pair(mr_ref, configured); + + // 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; } // 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 = *std::get<0>(make_host_mr(std::nullopt)); + static rmm::host_async_resource_ref mr_ref = make_host_mr(std::nullopt); return mr_ref; } @@ -280,8 +285,9 @@ rmm::host_async_resource_ref get_host_memory_resource() bool config_default_host_memory_resource(host_mr_options const& opts) { std::scoped_lock lock{host_mr_mutex()}; - auto configured = std::get<1>(make_host_mr(opts)); - return configured; + auto did_configure = false; + make_host_mr(opts, &did_configure); + return did_configure; } } // namespace cudf::io