Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Return boolean from config_host_memory_resource instead of throwing #15815

Merged
4 changes: 3 additions & 1 deletion cpp/include/cudf/io/memory_resource.hpp
Original file line number Diff line number Diff line change
Expand Up @@ -57,7 +57,9 @@ 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.
*/
void config_default_host_memory_resource(host_mr_options const& opts);
bool config_default_host_memory_resource(host_mr_options const& opts);

} // namespace cudf::io
20 changes: 13 additions & 7 deletions cpp/src/io/utilities/config_utils.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -244,16 +244,20 @@ 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<host_mr_options> const& opts)
CUDF_EXPORT rmm::host_async_resource_ref& make_host_mr(std::optional<host_mr_options> const& opts,
bool* did_configure = nullptr)
vuule marked this conversation as resolved.
Show resolved Hide resolved
{
static rmm::host_async_resource_ref* mr_ref = nullptr;
bool configured = false;
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");
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;
}

Expand All @@ -278,10 +282,12 @@ rmm::host_async_resource_ref get_host_memory_resource()
return host_mr();
}

void config_default_host_memory_resource(host_mr_options const& opts)
bool config_default_host_memory_resource(host_mr_options const& opts)
{
std::scoped_lock lock{host_mr_mutex()};
make_host_mr(opts);
auto did_configure = false;
make_host_mr(opts, &did_configure);
return did_configure;
}

} // namespace cudf::io
7 changes: 5 additions & 2 deletions java/src/main/java/ai/rapids/cudf/PinnedMemoryPool.java
Original file line number Diff line number Diff line change
Expand Up @@ -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);
}

}
5 changes: 4 additions & 1 deletion java/src/main/java/ai/rapids/cudf/Rmm.java
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
11 changes: 5 additions & 6 deletions java/src/main/native/src/RmmJni.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -1035,7 +1035,6 @@ 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
prior_cuio_host_mr() = cudf::io::set_host_memory_resource(*pinned_fallback_mr);
}
CATCH_STD(env, )
Expand Down Expand Up @@ -1107,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)
}
}
Loading