[Java] One PinnedMemoryBuffer per CuVSResourcesImpl (#1441)

While profiling cuvs-java, we found that allocating a `PinnedMemoryBuffer` for each host->device or device->host memory copy was unnecessary and wasteful.
This PR moves the allocation of a `PinnedMemoryBuffer` to `CuVSResourcesImpl`, so that the buffer can be cached and reused. Since `CuVSResources` are already meant to be per-thread, this is safe, as the `PinnedMemoryBuffer` will never be used concurrently.
In order to do it cleanly, we introduced two named `ScopedAccess` classes and a helper method that will always find its way to the internal `MemorySegment` used by native functions to access the buffer, without the need to expose it via the public interface.

Authors:
  - Lorenzo Dematté (https://github.com/ldematte)
  - Ben Frederickson (https://github.com/benfred)
  - MithunR (https://github.com/mythrocks)

Approvers:
  - MithunR (https://github.com/mythrocks)

URL: #1441

Author: ldematte

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

@@ -11,6 +11,7 @@
 import java.util.Random;
 
 @Fork(value = 1, warmups = 0)
+@Threads(1) // Sharing resources
 @State(Scope.Benchmark)
 public class CuVSDeviceMatrixBenchmarks {
 
@@ -82,15 +83,13 @@ public void matrixCopyDeviceToHost() {
 
   @Benchmark
   public void matrixDeviceBuilder() throws Throwable {
-    try (CuVSResources resources = CuVSResources.create()) {
-      var builder = CuVSMatrix.deviceBuilder(resources, size, dims, CuVSMatrix.DataType.FLOAT);
+    var builder = CuVSMatrix.deviceBuilder(resources, size, dims, CuVSMatrix.DataType.FLOAT);
 
-      for (int i = 0; i < size; ++i) {
-        var array = data[i];
-        builder.addVector(array);
-      }
-      CuVSDeviceMatrix matrix = builder.build();
-      matrix.close();
+    for (int i = 0; i < size; ++i) {
+      var array = data[i];
+      builder.addVector(array);
     }
+    CuVSDeviceMatrix matrix = builder.build();
+    matrix.close();
   }
 }

java/cuvs-java/src/main/java/com/nvidia/cuvs/CuVSMatrix.java

@@ -17,10 +17,20 @@
 public interface CuVSMatrix extends AutoCloseable {
 
   enum DataType {
-    FLOAT,
-    INT,
-    UINT,
-    BYTE
+    FLOAT(4),
+    INT(4),
+    UINT(4),
+    BYTE(1);
+
+    private final int bytes;
+
+    DataType(int bytes) {
+      this.bytes = bytes;
+    }
+
+    public int bytes() {
+      return bytes;
+    }
   }
 
   enum MemoryKind {

java/cuvs-java/src/main/java/com/nvidia/cuvs/DelegatingScopedAccess.java

@@ -0,0 +1,29 @@
+/*
+ * SPDX-FileCopyrightText: Copyright (c) 2025, NVIDIA CORPORATION.
+ * SPDX-License-Identifier: Apache-2.0
+ */
+package com.nvidia.cuvs;
+
+public final class DelegatingScopedAccess implements CuVSResources.ScopedAccess {
+  private final CuVSResources.ScopedAccess inner;
+  private final Runnable closeAction;
+
+  DelegatingScopedAccess(CuVSResources.ScopedAccess inner, Runnable closeAction) {
+    this.inner = inner;
+    this.closeAction = closeAction;
+  }
+
+  public CuVSResources.ScopedAccess inner() {
+    return inner;
+  }
+
+  @Override
+  public long handle() {
+    return inner.handle();
+  }
+
+  @Override
+  public void close() {
+    closeAction.run();
+  }
+}

java/cuvs-java/src/main/java/com/nvidia/cuvs/SynchronizedCuVSResources.java

@@ -27,17 +27,7 @@ static CuVSResources create() throws Throwable {
   @Override
   public ScopedAccess access() {
     lock.lock();
-    return new ScopedAccess() {
-      @Override
-      public long handle() {
-        return inner.access().handle();
-      }
-
-      @Override
-      public void close() {
-        lock.unlock();
-      }
-    };
+    return new DelegatingScopedAccess(inner.access(), lock::unlock);
   }
 
   @Override

java/cuvs-java/src/main/java22/com/nvidia/cuvs/internal/CuVSDeviceMatrixImpl.java

@@ -15,15 +15,55 @@
 
 public class CuVSDeviceMatrixImpl extends CuVSMatrixBaseImpl implements CuVSDeviceMatrix {
 
-  private long bufferedMatrixRowStart = 0;
-  private long bufferedMatrixRowEnd = 0;
+  private interface RowAccessStrategy {
+    RowView getRow(long row);
+  }
 
   private final CuVSResources resources;
 
   private final long rowStride;
   private final long columnStride;
+  private final long rowSize;
+  private final long valueByteSize;
+
+  private final RowAccessStrategy rowAccessStrategy;
 
-  private final PinnedMemoryBuffer hostBuffer;
+  private class BufferedRowAccessStrategy implements RowAccessStrategy {
+    private long bufferedMatrixRowStart = 0;
+    private long bufferedMatrixRowEnd = 0;
+
+    @Override
+    public RowView getRow(long row) {
+      try (var access = resources.access()) {
+        var hostBuffer = CuVSResourcesImpl.getHostBuffer(access);
+        long rowBytes = columns * valueByteSize;
+        if (row < bufferedMatrixRowStart || row >= bufferedMatrixRowEnd) {
+          var endRow = Math.min(row + (PinnedMemoryBuffer.CHUNK_BYTES / rowBytes), size);
+          populateBuffer(access, row, endRow, hostBuffer);
+          bufferedMatrixRowStart = row;
+          bufferedMatrixRowEnd = endRow;
+        }
+        var startRow = row - bufferedMatrixRowStart;
+        return new SliceRowView(
+            hostBuffer.asSlice(startRow * rowSize, rowBytes),
+            columns,
+            valueLayout,
+            dataType,
+            valueByteSize);
+      }
+    }
+  }
+
+  private class DirectRowAccessStrategy implements RowAccessStrategy {
+    @Override
+    public RowView getRow(long row) {
+      try (var access = resources.access()) {
+        var memorySegment = Arena.ofAuto().allocate(size * valueByteSize);
+        populateBuffer(access, row, row + 1, memorySegment);
+        return new SliceRowView(memorySegment, columns, valueLayout, dataType, valueByteSize);
+      }
+    }
+  }
 
   protected CuVSDeviceMatrixImpl(
       CuVSResources resources,
@@ -48,7 +88,15 @@ protected CuVSDeviceMatrixImpl(
     this.resources = resources;
     this.rowStride = rowStride;
     this.columnStride = columnStride;
-    this.hostBuffer = new PinnedMemoryBuffer(size, columns, valueLayout);
+
+    this.valueByteSize = valueLayout.byteSize();
+    this.rowSize = rowStride > 0 ? rowStride * valueByteSize : columns * valueByteSize;
+    if (rowSize > PinnedMemoryBuffer.CHUNK_BYTES) {
+      // The shared buffer is too small for this row size, use a direct access strategy
+      this.rowAccessStrategy = new DirectRowAccessStrategy();
+    } else {
+      this.rowAccessStrategy = new BufferedRowAccessStrategy();
+    }
   }
 
   @Override
@@ -63,15 +111,14 @@ public MemorySegment toTensor(Arena arena) {
         arena, memorySegment, new long[] {size, columns}, strides, code(), bits(), kDLCUDA());
   }
 
-  private void populateBuffer(long startRow) {
+  private void populateBuffer(
+      CuVSResources.ScopedAccess resourceAccess,
+      long startRow,
+      long endRow,
+      MemorySegment bufferAddress) {
     try (var localArena = Arena.ofConfined()) {
-      long rowBytes = columns * valueLayout.byteSize();
-      var endRow = Math.min(startRow + (hostBuffer.size() / rowBytes), size);
       var rowCount = endRow - startRow;
 
-      //      System.out.printf(
-      //          Locale.ROOT, "startRow: %d, endRow %d, count: %d\n", startRow, endRow, rowCount);
-
       MemorySegment sliceManagedTensor = DLManagedTensor.allocate(localArena);
       DLManagedTensor.dl_tensor(sliceManagedTensor, DLTensor.allocate(localArena));
 
@@ -85,40 +132,18 @@ private void populateBuffer(long startRow) {
 
       MemorySegment bufferTensor =
           prepareTensor(
-              localArena,
-              hostBuffer.address(),
-              new long[] {rowCount, columns},
-              code(),
-              bits(),
-              kDLCPU());
-
-      try (var resourceAccess = resources.access()) {
-        checkCuVSError(
-            cuvsMatrixCopy(resourceAccess.handle(), sliceManagedTensor, bufferTensor),
-            "cuvsMatrixCopy");
-        checkCuVSError(cuvsStreamSync(resourceAccess.handle()), "cuvsStreamSync");
-
-        bufferedMatrixRowStart = startRow;
-        bufferedMatrixRowEnd = endRow;
-      }
+              localArena, bufferAddress, new long[] {rowCount, columns}, code(), bits(), kDLCPU());
+
+      checkCuVSError(
+          cuvsMatrixCopy(resourceAccess.handle(), sliceManagedTensor, bufferTensor),
+          "cuvsMatrixCopy");
+      checkCuVSError(cuvsStreamSync(resourceAccess.handle()), "cuvsStreamSync");
     }
   }
 
   @Override
   public RowView getRow(long row) {
-    if (row < bufferedMatrixRowStart || row >= bufferedMatrixRowEnd) {
-      populateBuffer(row);
-    }
-    var valueByteSize = valueLayout.byteSize();
-    var startRow = row - bufferedMatrixRowStart;
-
-    var rowSize = rowStride > 0 ? rowStride * valueByteSize : columns * valueByteSize;
-    return new SliceRowView(
-        hostBuffer.address().asSlice(startRow * rowSize, columns * valueByteSize),
-        columns,
-        valueLayout,
-        dataType,
-        valueByteSize);
+    return rowAccessStrategy.getRow(row);
   }
 
   @Override
@@ -199,9 +224,7 @@ public void toDevice(CuVSDeviceMatrix targetMatrix, CuVSResources cuVSResources)
   }
 
   @Override
-  public void close() {
-    hostBuffer.close();
-  }
+  public void close() {}
 
   private static class CuVSDeviceMatrixDelegate implements CuVSDeviceMatrix, CuVSMatrixInternal {
     private final CuVSDeviceMatrixImpl deviceMatrix;

java/cuvs-java/src/main/java22/com/nvidia/cuvs/internal/CuVSResourcesImpl.java

@@ -9,7 +9,10 @@
 import static com.nvidia.cuvs.internal.panama.headers_h_1.C_INT;
 
 import com.nvidia.cuvs.CuVSResources;
+import com.nvidia.cuvs.DelegatingScopedAccess;
+import com.nvidia.cuvs.internal.common.PinnedMemoryBuffer;
 import java.lang.foreign.Arena;
+import java.lang.foreign.MemorySegment;
 import java.nio.file.Path;
 
 /**
@@ -24,6 +27,8 @@ public class CuVSResourcesImpl implements CuVSResources {
   private final ScopedAccess access;
   private final int deviceId;
 
+  private final PinnedMemoryBuffer hostBuffer = new PinnedMemoryBuffer();
+
   /**
    * Constructor that allocates the resources needed for cuVS
    *
@@ -37,16 +42,7 @@ public CuVSResourcesImpl(Path tempDirectory) {
       var deviceIdPtr = localArena.allocate(C_INT);
       checkCuVSError(cuvsDeviceIdGet(resourceHandle, deviceIdPtr), "cuvsDeviceIdGet");
       this.deviceId = deviceIdPtr.get(C_INT, 0);
-      this.access =
-          new ScopedAccess() {
-            @Override
-            public long handle() {
-              return resourceHandle;
-            }
-
-            @Override
-            public void close() {}
-          };
+      this.access = new ScopedAccessWithHostBuffer(resourceHandle, hostBuffer.address());
     }
   }
 
@@ -65,11 +61,25 @@ public void close() {
     synchronized (this) {
       int returnValue = cuvsResourcesDestroy(resourceHandle);
       checkCuVSError(returnValue, "cuvsResourcesDestroy");
+      hostBuffer.close();
     }
   }
 
   @Override
   public Path tempDirectory() {
     return tempDirectory;
   }
+
+  public static MemorySegment getHostBuffer(ScopedAccess access) {
+
+    while (access instanceof DelegatingScopedAccess delegatingScopedAccess) {
+      access = delegatingScopedAccess.inner();
+    }
+
+    if (access instanceof ScopedAccessWithHostBuffer withHostBuffer) {
+      return withHostBuffer.hostBuffer();
+    }
+
+    throw new IllegalArgumentException("Unsupported access type: " + access.getClass().getName());
+  }
 }

java/cuvs-java/src/main/java22/com/nvidia/cuvs/internal/ScopedAccessWithHostBuffer.java

@@ -0,0 +1,30 @@
+/*
+ * SPDX-FileCopyrightText: Copyright (c) 2025, NVIDIA CORPORATION.
+ * SPDX-License-Identifier: Apache-2.0
+ */
+package com.nvidia.cuvs.internal;
+
+import com.nvidia.cuvs.CuVSResources;
+import java.lang.foreign.MemorySegment;
+
+class ScopedAccessWithHostBuffer implements CuVSResources.ScopedAccess {
+  private final long resourceHandle;
+  private final MemorySegment hostBuffer;
+
+  public ScopedAccessWithHostBuffer(long resourceHandle, MemorySegment hostBuffer) {
+    this.resourceHandle = resourceHandle;
+    this.hostBuffer = hostBuffer;
+  }
+
+  @Override
+  public long handle() {
+    return resourceHandle;
+  }
+
+  public MemorySegment hostBuffer() {
+    return hostBuffer;
+  }
+
+  @Override
+  public void close() {}
+}