[Java] Uniform toHost/toDevice to work across all CuVSMatrix classes (#1328)

This PR moves `toHost`/`toDevice` operations to `CuVSMatrix` in order to uniform them; implementations then do not need to worry or perform checks on the input type, but simply accept any `CuVSMatrix` and then "request" the type they want or need (no ifs, custom cudaAlloc or cudaMemcpy operations needed):
```
CuVSDeviceMatrix matrix = input.toDevice(resources);
```
If the source matrix is a `CuVSHostMatrix`, device memory will be allocated and data will be copied; if the source matrix is a already a `CuVSDeviceMatrix`, a (wrapped) reference to the original data is returned (very fast, no copies, no additional allocations).

The wrapper is currently implemented to follow C++ `weak_ptr` semantics: it delegates everything to the wrapped matrix, the only difference is on `close()` (which is a no-op in the wrapper, as you'd expect). The lifecycle of the data follows the one of the original matrix, the "weak" delegate does not influence it in any way. This way the caller can handle it uniformly (e.g. with try-with-resources), without side effects on the original matrix (which lifecycle will be already handled by its owner).

An alternative would be to give this C++ `shared_ptr` semantics, making it reference counted. I think this would be more complex and not necessary, but let me know.

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

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

URL: #1328

Author: ldematte

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

@@ -21,16 +21,7 @@
 public interface CuVSDeviceMatrix extends CuVSMatrix {
 
   /**
-   * Fills the provided, pre-allocated host matrix with data from this device matrix.
-   * The content of the provided host matrix will be overwritten; the 2 matrices must have the
-   * same element type and dimension.
-   *
-   * @param hostMatrix  the host-memory-backed matrix to fill.
-   */
-  void toHost(CuVSHostMatrix hostMatrix);
-
-  /**
-   * Returns a new, host matrix with data from this device matrix.
+   * Returns a new host matrix with data from this device matrix.
    * The returned host matrix will need to be managed by the caller, which will be
    * responsible to call {@link CuVSMatrix#close()} to free its resources when done.
    */
@@ -39,4 +30,67 @@ default CuVSHostMatrix toHost() {
     toHost(hostMatrix);
     return hostMatrix;
   }
+
+  default CuVSDeviceMatrix toDevice(CuVSResources resources) {
+    return new CuVSDeviceMatrixDelegate(this);
+  }
+
+  class CuVSDeviceMatrixDelegate implements CuVSDeviceMatrix {
+
+    private final CuVSDeviceMatrix deviceMatrix;
+
+    private CuVSDeviceMatrixDelegate(CuVSDeviceMatrix deviceMatrix) {
+      this.deviceMatrix = deviceMatrix;
+    }
+
+    @Override
+    public long size() {
+      return deviceMatrix.size();
+    }
+
+    @Override
+    public long columns() {
+      return deviceMatrix.columns();
+    }
+
+    @Override
+    public DataType dataType() {
+      return deviceMatrix.dataType();
+    }
+
+    @Override
+    public RowView getRow(long row) {
+      return deviceMatrix.getRow(row);
+    }
+
+    @Override
+    public void toArray(int[][] array) {
+      deviceMatrix.toArray(array);
+    }
+
+    @Override
+    public void toArray(float[][] array) {
+      deviceMatrix.toArray(array);
+    }
+
+    @Override
+    public void toArray(byte[][] array) {
+      deviceMatrix.toArray(array);
+    }
+
+    @Override
+    public void toHost(CuVSHostMatrix hostMatrix) {
+      deviceMatrix.toHost(hostMatrix);
+    }
+
+    @Override
+    public void toDevice(CuVSDeviceMatrix deviceMatrix, CuVSResources cuVSResources) {
+      this.deviceMatrix.toDevice(deviceMatrix, cuVSResources);
+    }
+
+    @Override
+    public void close() {
+      // Do nothing
+    }
+  }
 }

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

@@ -20,4 +20,77 @@
  */
 public interface CuVSHostMatrix extends CuVSMatrix {
   int get(int row, int col);
+
+  default CuVSHostMatrix toHost() {
+    return new CuVSHostMatrixDelegate(this);
+  }
+
+  default CuVSDeviceMatrix toDevice(CuVSResources resources) {
+    var deviceMatrix = CuVSMatrix.deviceBuilder(resources, size(), columns(), dataType()).build();
+    toDevice(deviceMatrix, resources);
+    return deviceMatrix;
+  }
+
+  class CuVSHostMatrixDelegate implements CuVSHostMatrix {
+    private final CuVSHostMatrix hostMatrix;
+
+    public CuVSHostMatrixDelegate(CuVSHostMatrix cuVSHostMatrix) {
+      this.hostMatrix = cuVSHostMatrix;
+    }
+
+    @Override
+    public int get(int row, int col) {
+      return hostMatrix.get(row, col);
+    }
+
+    @Override
+    public long size() {
+      return hostMatrix.size();
+    }
+
+    @Override
+    public long columns() {
+      return hostMatrix.columns();
+    }
+
+    @Override
+    public DataType dataType() {
+      return hostMatrix.dataType();
+    }
+
+    @Override
+    public RowView getRow(long row) {
+      return hostMatrix.getRow(row);
+    }
+
+    @Override
+    public void toArray(int[][] array) {
+      hostMatrix.toArray(array);
+    }
+
+    @Override
+    public void toArray(float[][] array) {
+      hostMatrix.toArray(array);
+    }
+
+    @Override
+    public void toArray(byte[][] array) {
+      hostMatrix.toArray(array);
+    }
+
+    @Override
+    public void toHost(CuVSHostMatrix hostMatrix) {
+      this.hostMatrix.toHost(hostMatrix);
+    }
+
+    @Override
+    public void toDevice(CuVSDeviceMatrix deviceMatrix, CuVSResources cuVSResources) {
+      hostMatrix.toDevice(deviceMatrix, cuVSResources);
+    }
+
+    @Override
+    public void close() {
+      // Do nothing
+    }
+  }
 }

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

@@ -177,6 +177,42 @@ static Builder<CuVSDeviceMatrix> deviceBuilder(
    */
   void toArray(byte[][] array);
 
+  /**
+   * Fills the provided, pre-allocated host matrix with data from this matrix.
+   * The content of the provided host matrix will be overwritten; the 2 matrices must have the
+   * same element type and dimension.
+   *
+   * @param hostMatrix  the host-memory-backed matrix to fill.
+   */
+  void toHost(CuVSHostMatrix hostMatrix);
+
+  /**
+   * Returns a host matrix; if the matrix is already a host matrix, a "weak" reference to the same host memory
+   * is returned. If the matrix is a device matrix, a newly allocated matrix will be populated with data from
+   * the device matrix.
+   * The returned host matrix will need to be managed by the caller, which will be
+   * responsible to call {@link CuVSMatrix#close()} to free its resources when done.
+   */
+  CuVSHostMatrix toHost();
+
+  /**
+   * Fills the provided, pre-allocated device matrix with data from this matrix.
+   * The content of the provided device matrix will be overwritten; the 2 matrices must have the
+   * same element type and dimension.
+   *
+   * @param deviceMatrix  the device-memory-backed matrix to fill.
+   */
+  void toDevice(CuVSDeviceMatrix deviceMatrix, CuVSResources cuVSResources);
+
+  /**
+   * Returns a device matrix; if this matrix is already a device matrix, a "weak" reference to the same host memory
+   * is returned. If the matrix is a host matrix, a newly allocated matrix will be populated with data from
+   * the host matrix.
+   * The returned device matrix will need to be managed by the caller, which will be
+   * responsible to call {@link CuVSMatrix#close()} to free its resources when done.
+   */
+  CuVSDeviceMatrix toDevice(CuVSResources cuVSResources);
+
   @Override
   void close();
 }

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

@@ -236,6 +236,11 @@ public void toHost(CuVSHostMatrix hostMatrix) {
     }
   }
 
+  @Override
+  public void toDevice(CuVSDeviceMatrix targetMatrix, CuVSResources cuVSResources) {
+    copyMatrix(this, (CuVSMatrixBaseImpl) targetMatrix, cuVSResources);
+  }
+
   @Override
   public void close() {
     if (hostBuffer != MemorySegment.NULL) {

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

@@ -18,7 +18,9 @@
 import static com.nvidia.cuvs.internal.common.Util.prepareTensor;
 import static com.nvidia.cuvs.internal.panama.headers_h.*;
 
+import com.nvidia.cuvs.CuVSDeviceMatrix;
 import com.nvidia.cuvs.CuVSHostMatrix;
+import com.nvidia.cuvs.CuVSResources;
 import com.nvidia.cuvs.RowView;
 import java.lang.foreign.*;
 import java.lang.invoke.VarHandle;
@@ -142,6 +144,19 @@ public void toArray(byte[][] array) {
     }
   }
 
+  @Override
+  public void toHost(CuVSHostMatrix hostMatrix) {
+    var targetMatrix = (CuVSHostMatrixImpl) hostMatrix;
+    var valueByteSize = valueLayout.byteSize();
+    MemorySegment.copy(
+        this.memorySegment, 0L, targetMatrix.memorySegment, 0L, size * columns * valueByteSize);
+  }
+
+  @Override
+  public void toDevice(CuVSDeviceMatrix deviceMatrix, CuVSResources cuVSResources) {
+    copyMatrix(this, (CuVSMatrixBaseImpl) deviceMatrix, cuVSResources);
+  }
+
   @Override
   public void close() {}
 

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

@@ -18,6 +18,7 @@
 import static com.nvidia.cuvs.internal.common.LinkerHelper.C_CHAR;
 import static com.nvidia.cuvs.internal.common.LinkerHelper.C_FLOAT;
 import static com.nvidia.cuvs.internal.common.LinkerHelper.C_INT;
+import static com.nvidia.cuvs.internal.common.Util.checkCuVSError;
 import static com.nvidia.cuvs.internal.panama.headers_h.*;
 
 import com.nvidia.cuvs.CuVSMatrix;
@@ -49,6 +50,29 @@ protected CuVSMatrixBaseImpl(
     this.columns = columns;
   }
 
+  protected static void copyMatrix(
+      CuVSMatrixBaseImpl sourceMatrix, CuVSMatrixBaseImpl targetMatrix, CuVSResources resources) {
+    if (targetMatrix.columns() != sourceMatrix.columns
+        || targetMatrix.size() != sourceMatrix.size) {
+      throw new IllegalArgumentException(
+          "Source and target matrices must have the same dimensions");
+    }
+    if (targetMatrix.dataType() != sourceMatrix.dataType) {
+      throw new IllegalArgumentException("Source and target matrices must have the same dataType");
+    }
+
+    try (var localArena = Arena.ofConfined()) {
+      var targetTensor = targetMatrix.toTensor(localArena);
+
+      try (var resourceAccess = resources.access()) {
+        var cuvsRes = resourceAccess.handle();
+        var sourceTensor = sourceMatrix.toTensor(localArena);
+        checkCuVSError(cuvsMatrixCopy(cuvsRes, sourceTensor, targetTensor), "cuvsMatrixCopy");
+        checkCuVSError(cuvsStreamSync(cuvsRes), "cuvsStreamSync");
+      }
+    }
+  }
+
   @Override
   public long size() {
     return size;