Merge branch 'branch-25.10' into java/device-matrix-cagra-index-tests

Author: mythrocks

.github/workflows/build.yaml

@@ -51,8 +51,8 @@ jobs:
       fail-fast: false
       matrix:
         cuda_version:
-          - '12.9.1'
-          - '13.0.0'
+          - &latest_cuda12 '12.9.1'
+          - &latest_cuda13 '13.0.1'
     with:
       build_type: ${{ inputs.build_type || 'branch' }}
       branch: ${{ inputs.branch }}
@@ -72,8 +72,8 @@ jobs:
       fail-fast: false
       matrix:
         cuda_version:
-          - '12.9.1'
-          - '13.0.0'
+          - *latest_cuda12
+          - *latest_cuda13
     with:
       build_type: ${{ inputs.build_type || 'branch' }}
       branch: ${{ inputs.branch }}
@@ -93,8 +93,8 @@ jobs:
       fail-fast: false
       matrix:
         cuda_version:
-          - '12.9.1'
-          - '13.0.0'
+          - *latest_cuda12
+          - *latest_cuda13
     with:
       build_type: ${{ inputs.build_type || 'branch' }}
       branch: ${{ inputs.branch }}

.github/workflows/pr.yaml

@@ -159,8 +159,8 @@ jobs:
       fail-fast: false
       matrix:
         cuda_version:
-          - '12.9.1'
-          - '13.0.0'
+          - &latest_cuda12 '12.9.1'
+          - &latest_cuda13 '13.0.1'
     with:
       build_type: pull-request
       node_type: "gpu-l4-latest-1"
@@ -189,8 +189,8 @@ jobs:
       fail-fast: false
       matrix:
         cuda_version:
-          - '12.9.1'
-          - '13.0.0'
+          - *latest_cuda12
+          - *latest_cuda13
     with:
       build_type: pull-request
       node_type: "gpu-l4-latest-1"
@@ -207,8 +207,8 @@ jobs:
       fail-fast: false
       matrix:
         cuda_version:
-          - '12.9.1'
-          - '13.0.0'
+          - *latest_cuda12
+          - *latest_cuda13
     with:
       build_type: pull-request
       node_type: "gpu-l4-latest-1"

.github/workflows/test.yaml

@@ -60,7 +60,7 @@ jobs:
       matrix:
         cuda_version:
           - '12.9.1'
-          - '13.0.0'
+          - '13.0.1'
     with:
       build_type: ${{ inputs.build_type }}
       branch: ${{ inputs.branch }}

cpp/src/cluster/detail/kmeans_mg.cuh

@@ -175,9 +175,20 @@ void initKMeansPlusPlus(const raft::resources& handle,
   //          X which will be used as the initial centroid for kmeans++
   //    1.3 - Communicate the initial centroid chosen by rank-r' to all other
   //          ranks
-  std::mt19937 gen(params.rng_state.seed);
-  std::uniform_int_distribution<> dis(0, n_rank - 1);
-  int rp = dis(gen);
+  // Choose rp on rank 0 and broadcast to all ranks to guarantee agreement
+  int rp = 0;
+  if (my_rank == KMEANS_COMM_ROOT) {
+    std::mt19937 gen(params.rng_state.seed);
+    std::uniform_int_distribution<> dis(0, n_rank - 1);
+    rp = dis(gen);
+  }
+  {
+    rmm::device_scalar<int> rp_d(stream);
+    raft::copy(rp_d.data(), &rp, 1, stream);
+    comm.bcast<int>(rp_d.data(), 1, /*root=*/KMEANS_COMM_ROOT, stream);
+    raft::copy(&rp, rp_d.data(), 1, stream);
+    raft::resource::sync_stream(handle);
+  }
 
   // buffer to flag the sample that is chosen as initial centroids
   std::vector<std::uint8_t> h_isSampleCentroid(n_samples);

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

@@ -30,67 +30,4 @@ 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

@@ -21,76 +21,9 @@
 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/java22/com/nvidia/cuvs/internal/BruteForceIndexImpl.java

@@ -78,8 +78,8 @@ private BruteForceIndexImpl(
     Objects.requireNonNull(dataset);
     try (dataset) {
       this.resources = resources;
-      assert dataset instanceof CuVSMatrixBaseImpl;
-      this.bruteForceIndexReference = build((CuVSMatrixBaseImpl) dataset, bruteForceIndexParams);
+      assert dataset instanceof CuVSMatrixInternal;
+      this.bruteForceIndexReference = build((CuVSMatrixInternal) dataset, bruteForceIndexParams);
     }
   }
 
@@ -124,7 +124,7 @@ public void close() {
    * index
    */
   private IndexReference build(
-      CuVSMatrixBaseImpl dataset, BruteForceIndexParams bruteForceIndexParams) {
+      CuVSMatrixInternal dataset, BruteForceIndexParams bruteForceIndexParams) {
     long rows = dataset.size();
     long cols = dataset.columns();
 

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

@@ -77,8 +77,8 @@ private CagraIndexImpl(
       CagraIndexParams indexParameters, CuVSMatrix dataset, CuVSResources resources) {
     Objects.requireNonNull(dataset);
     this.resources = resources;
-    assert dataset instanceof CuVSMatrixBaseImpl;
-    this.cagraIndexReference = build(indexParameters, (CuVSMatrixBaseImpl) dataset);
+    assert dataset instanceof CuVSMatrixInternal;
+    this.cagraIndexReference = build(indexParameters, (CuVSMatrixInternal) dataset);
   }
 
   /**
@@ -123,11 +123,11 @@ private CagraIndexImpl(
 
     this.resources = resources;
 
-    assert graph instanceof CuVSMatrixBaseImpl;
-    assert dataset instanceof CuVSMatrixBaseImpl;
+    assert graph instanceof CuVSMatrixInternal;
+    assert dataset instanceof CuVSMatrixInternal;
 
     this.cagraIndexReference =
-        fromGraph(metric, (CuVSMatrixBaseImpl) graph, (CuVSMatrixBaseImpl) dataset);
+        fromGraph(metric, (CuVSMatrixInternal) graph, (CuVSMatrixInternal) dataset);
   }
 
   private void checkNotDestroyed() {
@@ -160,7 +160,7 @@ public void close() {
    * @return an instance of {@link IndexReference} that holds the pointer to the
    *         index
    */
-  private IndexReference build(CagraIndexParams indexParameters, CuVSMatrixBaseImpl dataset) {
+  private IndexReference build(CagraIndexParams indexParameters, CuVSMatrixInternal dataset) {
     long rows = dataset.size();
 
     try (var indexParams = segmentFromIndexParams(indexParameters);
@@ -409,8 +409,8 @@ public CuVSDeviceMatrix getGraph() {
 
   private IndexReference fromGraph(
       CagraIndexParams.CuvsDistanceType metric,
-      CuVSMatrixBaseImpl graph,
-      CuVSMatrixBaseImpl dataset) {
+      CuVSMatrixInternal graph,
+      CuVSMatrixInternal dataset) {
     try (var localArena = Arena.ofConfined()) {
       var index = createCagraIndex();
       try (var resourcesAccess = resources.access()) {