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 68453c924d6..9c015fee409 100644 --- a/java/src/main/native/src/RmmJni.cpp +++ b/java/src/main/native/src/RmmJni.cpp @@ -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, ) +} }