Skip to content

Commit

Permalink
Revert "Revert "[JNI] Expose java API for cudf::io::config_host_memor…
Browse files Browse the repository at this point in the history
…y_resource (rapidsai#15745)""

This reverts commit ff392b5.
  • Loading branch information
abellina committed May 22, 2024
1 parent fd0d2b4 commit 8365382
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 @@ -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, )
}
}

0 comments on commit 8365382

Please sign in to comment.