[Java] Bindings, tests and benchmarks for RMM pooled memory

java/benchmarks/src/main/java/com/nvidia/cuvs/CuVSDeviceMatrixBenchmarks.java

@@ -9,6 +9,7 @@
 import org.openjdk.jmh.infra.Blackhole;
 
 import java.util.Random;
+import java.util.concurrent.*;
 
 @Fork(value = 1, warmups = 0)
 @State(Scope.Benchmark)
@@ -21,6 +22,8 @@ public class CuVSDeviceMatrixBenchmarks {
   private int size;
 
   private static final Random random = new Random();
+  private static final int ALLOCATIONS = 10_000;
+  private static final int NUM_THREADS = 10;
 
   private float[][] data;
 
@@ -93,4 +96,46 @@ public void matrixDeviceBuilder() throws Throwable {
       matrix.close();
     }
   }
+
+  @Benchmark
+  @Threads(1)
+  public void multipleDeviceAllocations() throws Throwable {
+    try (CuVSResources resources = CuVSResources.create()) {
+      for (int i = 0; i < ALLOCATIONS; ++i) {
+        var builder = CuVSMatrix.deviceBuilder(resources, size, dims, CuVSMatrix.DataType.FLOAT);
+        var matrix = builder.build();
+        matrix.close();
+      }
+    }
+  }
+
+  @Benchmark
+  @Threads(1)
+  public void multipleConcurrentDeviceAllocations() throws Throwable {
+    final var allocationsPerThread = ALLOCATIONS / NUM_THREADS;
+    Utils.runConcurrently(false, NUM_THREADS, () -> {
+      try (CuVSResources resources = CuVSResources.create()) {
+        for (int i = 0; i < allocationsPerThread; ++i) {
+          var builder = CuVSMatrix.deviceBuilder(resources, size, dims, CuVSMatrix.DataType.FLOAT);
+          var matrix = builder.build();
+          matrix.close();
+        }
+      }
+    });
+  }
+
+  @Benchmark
+  @Threads(1)
+  public void multipleConcurrentDeviceAllocationsWithPooledMemory() throws Throwable {
+    final var allocationsPerThread = ALLOCATIONS / NUM_THREADS;
+    Utils.runConcurrently(true, NUM_THREADS, () -> {
+      try (CuVSResources resources = CuVSResources.create()) {
+        for (int i = 0; i < allocationsPerThread; ++i) {
+          var builder = CuVSMatrix.deviceBuilder(resources, size, dims, CuVSMatrix.DataType.FLOAT);
+          var matrix = builder.build();
+          matrix.close();
+        }
+      }
+    });
+  }
 }

java/benchmarks/src/main/java/com/nvidia/cuvs/Utils.java

@@ -0,0 +1,53 @@
+package com.nvidia.cuvs;
+
+import com.nvidia.cuvs.spi.CuVSProvider;
+
+import java.util.concurrent.*;
+import java.util.function.Supplier;
+
+class Utils {
+
+  @FunctionalInterface
+  interface TestFunction {
+    void apply() throws Throwable;
+  }
+
+  static void runConcurrently(boolean usePooledMemory, int nThreads, TestFunction testFunction)
+          throws ExecutionException, InterruptedException, TimeoutException {
+    try (ExecutorService parallelExecutor = Executors.newFixedThreadPool(nThreads)) {
+      if (usePooledMemory) {
+        CuVSProvider.provider().enableRMMPooledMemory(10, 60);
+      }
+      var futures = new CompletableFuture[nThreads];
+      for (int j = 0; j < nThreads; j++) {
+        futures[j] = CompletableFuture.runAsync(() -> {
+          try {
+            testFunction.apply();
+          } catch (Throwable e) {
+            throw new RuntimeException(e);
+          }
+        }, parallelExecutor);
+      }
+
+      CompletableFuture.allOf(futures)
+              .exceptionally(Utils::fail)
+              .get(2000, TimeUnit.SECONDS);
+    } finally {
+      if (usePooledMemory) {
+        CuVSProvider.provider().resetRMMPooledMemory();
+      }
+    }
+  }
+
+  static Throwable unwrap(Throwable t) {
+    var root = t;
+    while (root.getCause() != null) {
+      root = root.getCause();
+    }
+    return root;
+  }
+
+  private static Void fail(Throwable t) {
+    throw new AssertionError("Exception while executing: " + unwrap(t));
+  }
+}

java/cuvs-java/src/main/java/com/nvidia/cuvs/spi/CuVSProvider.java

@@ -157,6 +157,29 @@ default CagraIndex mergeCagraIndexes(CagraIndex[] indexes, CagraMergeParams merg
 
   java.util.logging.Level getLogLevel();
 
+  /**
+   * Switch RMM allocations (used internally by various cuVS algorithms and by the default implementation of
+   * {@link CuVSDeviceMatrix}) to use pooled memory.
+   * This operation has a global effect, and will affect all resources on the current device.
+   *
+   * @param initialPoolSizePercent The initial pool size, in percentage of the total GPU memory
+   * @param maxPoolSizePercent The maximum pool size, in percentage of the total GPU memory
+   */
+  void enableRMMPooledMemory(int initialPoolSizePercent, int maxPoolSizePercent);
+
+  /**
+   * Switch RMM allocations (used internally by various cuVS algorithms and by the default implementation of
+   * {@link CuVSDeviceMatrix}) to use pooled memory.
+   * This operation has a global effect, and will affect all resources on the current device.
+   *
+   * @param initialPoolSizePercent The initial pool size, in percentage of the total GPU memory
+   * @param maxPoolSizePercent The maximum pool size, in percentage of the total GPU memory
+   */
+  void enableRMMManagedPooledMemory(int initialPoolSizePercent, int maxPoolSizePercent);
+
+  /** Disables pooled memory on the current device, reverting back to the default setting.  */
+  void resetRMMPooledMemory();
+
   /** Retrieves the system-wide provider. */
   static CuVSProvider provider() {
     return CuVSServiceProvider.Holder.INSTANCE;

java/cuvs-java/src/main/java/com/nvidia/cuvs/spi/UnsupportedProvider.java

@@ -100,6 +100,21 @@ public Level getLogLevel() {
     throw new UnsupportedOperationException(reasons);
   }
 
+  @Override
+  public void enableRMMPooledMemory(int initialPoolSizePercent, int maxPoolSizePercent) {
+    throw new UnsupportedOperationException(reasons);
+  }
+
+  @Override
+  public void enableRMMManagedPooledMemory(int initialPoolSizePercent, int maxPoolSizePercent) {
+    throw new UnsupportedOperationException(reasons);
+  }
+
+  @Override
+  public void resetRMMPooledMemory() {
+    throw new UnsupportedOperationException(reasons);
+  }
+
   @Override
   public CuVSMatrix.Builder<CuVSDeviceMatrix> newDeviceMatrixBuilder(
       CuVSResources cuVSResources,

java/cuvs-java/src/main/java22/com/nvidia/cuvs/spi/JDKProvider.java

@@ -68,6 +68,9 @@ final class JDKProvider implements CuVSProvider {
     }
   }
 
+  private final cuvsRMMMemoryResourceReset cuvsRMMMemoryResourceResetInvoker =
+      cuvsRMMMemoryResourceReset.makeInvoker();
+
   private final cuvsGetLogLevel GET_LOG_LEVEL_INVOKER = cuvsGetLogLevel.makeInvoker();
 
   private JDKProvider() {}
@@ -338,6 +341,25 @@ public Level getLogLevel() {
     throw new IllegalArgumentException("Unexpected log level [" + logLevel + "]");
   }
 
+  @Override
+  public void enableRMMPooledMemory(int initialPoolSizePercent, int maxPoolSizePercent) {
+    checkCuVSError(
+        cuvsRMMPoolMemoryResourceEnable(initialPoolSizePercent, maxPoolSizePercent, false),
+        "cuvsRMMPoolMemoryResourceEnable");
+  }
+
+  @Override
+  public void enableRMMManagedPooledMemory(int initialPoolSizePercent, int maxPoolSizePercent) {
+    checkCuVSError(
+        cuvsRMMPoolMemoryResourceEnable(initialPoolSizePercent, maxPoolSizePercent, true),
+        "cuvsRMMPoolMemoryResourceEnable");
+  }
+
+  @Override
+  public void resetRMMPooledMemory() {
+    checkCuVSError(cuvsRMMMemoryResourceResetInvoker.apply(), "cuvsRMMMemoryResourceReset");
+  }
+
   @Override
   public CuVSMatrix.Builder<CuVSHostMatrix> newHostMatrixBuilder(
       long size, long columns, CuVSMatrix.DataType dataType) {

java/cuvs-java/src/test/java/com/nvidia/cuvs/BruteForceAndSearchIT.java

@@ -6,6 +6,7 @@
 
 import static com.carrotsearch.randomizedtesting.RandomizedTest.assumeTrue;
 
+import com.nvidia.cuvs.spi.CuVSProvider;
 import java.io.*;
 import java.nio.file.Files;
 import java.nio.file.Path;
@@ -14,6 +15,7 @@
 import java.util.Map;
 import java.util.UUID;
 import java.util.function.LongToIntFunction;
+import org.junit.After;
 import org.junit.Before;
 import org.junit.Test;
 
@@ -22,6 +24,12 @@ public class BruteForceAndSearchIT extends CuVSTestCase {
   @Before
   public void setup() {
     assumeTrue("not supported on " + System.getProperty("os.name"), isLinuxAmd64());
+    CuVSProvider.provider().enableRMMPooledMemory(10, 60);
+  }
+
+  @After
+  public void cleanup() {
+    CuVSProvider.provider().resetRMMPooledMemory();
   }
 
   // Sample data and query

java/cuvs-java/src/test/java/com/nvidia/cuvs/BruteForceRandomizedIT.java

@@ -7,9 +7,11 @@
 import static com.carrotsearch.randomizedtesting.RandomizedTest.assumeTrue;
 
 import com.carrotsearch.randomizedtesting.RandomizedRunner;
+import com.nvidia.cuvs.spi.CuVSProvider;
 import java.lang.invoke.MethodHandles;
 import java.util.BitSet;
 import java.util.List;
+import org.junit.After;
 import org.junit.Before;
 import org.junit.Test;
 import org.junit.runner.RunWith;
@@ -26,6 +28,12 @@ public void setup() {
     assumeTrue(isLinuxAmd64());
     initializeRandom();
     log.trace("Random context initialized for test.");
+    CuVSProvider.provider().enableRMMPooledMemory(10, 60);
+  }
+
+  @After
+  public void cleanup() {
+    CuVSProvider.provider().resetRMMPooledMemory();
   }
 
   @Test