Forward-merge branch-24.06 into branch-24.08

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);
+  }
+
 }

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.

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, )
+}
 }