From 8b7245548c63d1ce84031a0bd187cbfb8e072f8c Mon Sep 17 00:00:00 2001
From: Alessandro Bellina <abellina@nvidia.com>
Date: Mon, 20 May 2024 20:52:34 -0500
Subject: [PATCH] [JNI] Expose java API for
 cudf::io::config_host_memory_resource (#15745)

This PR depends on https://github.com/rapidsai/cudf/pull/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 #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: https://github.com/rapidsai/cudf/pull/15745
---
 .../main/java/ai/rapids/cudf/PinnedMemoryPool.java  | 13 +++++++++++++
 java/src/main/java/ai/rapids/cudf/Rmm.java          | 10 ++++++++++
 java/src/main/native/src/RmmJni.cpp                 | 11 +++++++++++
 3 files changed, 34 insertions(+)

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