Skip to content

Commit

Permalink
[JNI] Expose java API for cudf::io::config_host_memory_resource (rapi…
Browse files Browse the repository at this point in the history
…dsai#15745)

This PR depends on rapidsai#15665 and so it won't build until that PR merges.

Adds support for `cudf::io::config_host_memory_resource` which is being worked on in rapidsai#15665.  In 24.06 we are going to disable the cuDF pinned pool and look into this more in 24.08. 

We currently have a pinned pooled resource that has been setup to share pinned memory with other APIs we use from java, so we wanted to prevent extra pinned memory being created by default, and @vuule has added an API for us to call to accomplish this.

Authors:
  - Alessandro Bellina (https://github.com/abellina)
  - GALI PREM SAGAR (https://github.com/galipremsagar)

Approvers:
  - Nghia Truong (https://github.com/ttnghia)

URL: rapidsai#15745
  • Loading branch information
abellina committed May 21, 2024
1 parent 70922d0 commit 8b72455
Show file tree
Hide file tree
Showing 3 changed files with 34 additions and 0 deletions.
13 changes: 13 additions & 0 deletions java/src/main/java/ai/rapids/cudf/PinnedMemoryPool.java
Original file line number Diff line number Diff line change
Expand Up @@ -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);
}

}
10 changes: 10 additions & 0 deletions java/src/main/java/ai/rapids/cudf/Rmm.java
Original file line number Diff line number Diff line change
Expand Up @@ -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.
Expand Down
11 changes: 11 additions & 0 deletions java/src/main/native/src/RmmJni.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -1106,4 +1106,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, )
}
}

0 comments on commit 8b72455

Please sign in to comment.