[Java][C] Expose GPUInfo

cpp/include/cuvs/core/c_api.h

@@ -86,7 +86,7 @@ cuvsError_t cuvsResourcesDestroy(cuvsResources_t res);
 cuvsError_t cuvsStreamSet(cuvsResources_t res, cudaStream_t stream);
 
 /**
- * @brief Get the cudaStream_t from a cuvsResources_t t
+ * @brief Get the cudaStream_t from a cuvsResources_t
  *
  * @param[in] res cuvsResources_t opaque C handle
  * @param[out] stream cudaStream_t stream to queue CUDA kernels
@@ -101,6 +101,15 @@ cuvsError_t cuvsStreamGet(cuvsResources_t res, cudaStream_t* stream);
  * @return cuvsError_t
  */
 cuvsError_t cuvsStreamSync(cuvsResources_t res);
+
+/**
+ * @brief Get the id of the device associated with this cuvsResources_t
+ *
+ * @param[in] res cuvsResources_t opaque C handle
+ * @param[out] device_id int the id of the device associated with res
+ * @return cuvsError_t
+ */
+cuvsError_t cuvsDeviceIdGet(cuvsResources_t res, int* device_id);
 /** @} */
 
 /**

cpp/src/core/c_api.cpp

@@ -19,6 +19,8 @@
 #include <cuvs/version_config.h>
 
 #include <raft/core/resource/cuda_stream.hpp>
+#include <raft/core/resource/device_id.hpp>
+#include <raft/core/resource/resource_types.hpp>
 #include <raft/core/resources.hpp>
 #include <raft/util/cudart_utils.hpp>
 #include <rmm/cuda_stream_view.hpp>
@@ -73,6 +75,14 @@ extern "C" cuvsError_t cuvsStreamSync(cuvsResources_t res)
   });
 }
 
+extern "C" cuvsError_t cuvsDeviceIdGet(cuvsResources_t res, int* device_id)
+{
+  return cuvs::core::translate_exceptions([=] {
+    auto res_ptr = reinterpret_cast<raft::resources*>(res);
+    *device_id   = raft::resource::get_device_id(*res_ptr);
+  });
+}
+
 extern "C" cuvsError_t cuvsRMMAlloc(cuvsResources_t res, void** ptr, size_t bytes)
 {
   return cuvs::core::translate_exceptions([=] {

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

@@ -49,6 +49,13 @@ interface ScopedAccess extends AutoCloseable {
    */
   ScopedAccess access();
 
+  /**
+   * Get the logical id of the device associated with this resources object.
+   * Information about the device id is immutable, so it is safe to expose it without getting {@link ScopedAccess}
+   * to the enclosing resources.
+   */
+  int deviceId();
+
   /**
    * Closes this CuVSResources object and releases any resources associated with it.
    */

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

@@ -0,0 +1,26 @@
+/*
+ * Copyright (c) 2025, NVIDIA CORPORATION.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package com.nvidia.cuvs;
+
+/**
+ * Contains performance-related information associated to a {@link CuVSResources} and its GPU.
+ * Can be extended to report different types of GPU memory linked to the resources,
+ * e.g. the type and capacity of the underlying RMM {@code device_memory_resource}
+ *
+ * @param freeDeviceMemoryInBytes   free memory in bytes, as reported by the device driver
+ * @param totalDeviceMemoryInBytes  total device memory in bytes
+ */
+public record CuVSResourcesInfo(long freeDeviceMemoryInBytes, long totalDeviceMemoryInBytes) {}

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

@@ -18,11 +18,21 @@
 /**
  * Contains GPU information
  *
- * @param gpuId             id of the GPU starting from 0
- * @param name              ASCII string identifying device
- * @param freeMemory        returned free memory in bytes
- * @param totalMemory       returned total memory in bytes
- * @param computeCapability the compute capability of the device
+ * @param gpuId                     id of the GPU starting from 0
+ * @param name                      ASCII string identifying device
+ * @param totalDeviceMemoryInBytes  total device memory in bytes
+ * @param computeCapabilityMajor    the compute capability of the device (major)
+ * @param computeCapabilityMinor    the compute capability of the device (minor)
+ * @param supportsConcurrentCopy    whether the device can concurrently copy memory between host and device while
+ *                                  executing a kernel
+ * @param supportsConcurrentKernels whether the device supports executing multiple kernels within the same context
+ *                                  simultaneously
  */
 public record GPUInfo(
-    int gpuId, String name, long freeMemory, long totalMemory, float computeCapability) {}
+    int gpuId,
+    String name,
+    long totalDeviceMemoryInBytes,
+    int computeCapabilityMajor,
+    int computeCapabilityMinor,
+    boolean supportsConcurrentCopy,
+    boolean supportsConcurrentKernels) {}

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

@@ -0,0 +1,48 @@
+/*
+ * Copyright (c) 2025, NVIDIA CORPORATION.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package com.nvidia.cuvs;
+
+import java.util.List;
+
+public interface GPUInfoProvider {
+
+  int MIN_COMPUTE_CAPABILITY_MAJOR = 7;
+  int MIN_COMPUTE_CAPABILITY_MINOR = 0;
+
+  int MIN_DEVICE_MEMORY_IN_MB = 8192;
+
+  /**
+   * Gets all the available GPUs
+   *
+   * @return a list of {@link GPUInfo} objects with GPU details
+   */
+  List<GPUInfo> availableGPUs();
+
+  /**
+   * Get the list of compatible GPUs based on compute capability >= 7.0 and total
+   * memory >= 8GB
+   *
+   * @return a list of compatible GPUs. See {@link GPUInfo}
+   */
+  List<GPUInfo> compatibleGPUs();
+
+  /**
+   * Gets memory information relative to a {@link CuVSResources}
+   * @param resources from which to obtain memory information
+   * @return a {@link CuVSResourcesInfo} record containing the memory information
+   */
+  CuVSResourcesInfo getCurrentInfo(CuVSResources resources);
+}

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

@@ -51,6 +51,11 @@ public void close() {
     };
   }
 
+  @Override
+  public int deviceId() {
+    return inner.deviceId();
+  }
+
   @Override
   public void close() {
     inner.close();

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

@@ -128,6 +128,9 @@ default CagraIndex mergeCagraIndexes(CagraIndex[] indexes, CagraMergeParams merg
     return mergeCagraIndexes(indexes);
   }
 
+  /** Returns a {@link GPUInfoProvider} to query the system for GPU related information */
+  GPUInfoProvider gpuInfoProvider();
+
   /** Retrieves the system-wide provider. */
   static CuVSProvider provider() {
     return CuVSServiceProvider.Holder.INSTANCE;

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

@@ -66,6 +66,11 @@ public CuVSMatrix.Builder<CuVSDeviceMatrix> newDeviceMatrixBuilder(
     throw new UnsupportedOperationException();
   }
 
+  @Override
+  public GPUInfoProvider gpuInfoProvider() {
+    throw new UnsupportedOperationException();
+  }
+
   @Override
   public CuVSMatrix.Builder<CuVSDeviceMatrix> newDeviceMatrixBuilder(
       CuVSResources cuVSResources,

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

@@ -16,9 +16,8 @@
 package com.nvidia.cuvs.internal;
 
 import static com.nvidia.cuvs.internal.common.Util.checkCuVSError;
-import static com.nvidia.cuvs.internal.panama.headers_h.cuvsResourcesCreate;
-import static com.nvidia.cuvs.internal.panama.headers_h.cuvsResourcesDestroy;
-import static com.nvidia.cuvs.internal.panama.headers_h.cuvsResources_t;
+import static com.nvidia.cuvs.internal.panama.headers_h.*;
+import static com.nvidia.cuvs.internal.panama.headers_h_1.C_INT;
 
 import com.nvidia.cuvs.CuVSResources;
 import java.lang.foreign.Arena;
@@ -34,6 +33,7 @@ public class CuVSResourcesImpl implements CuVSResources {
   private final Path tempDirectory;
   private final long resourceHandle;
   private final ScopedAccess access;
+  private final int deviceId;
 
   /**
    * Constructor that allocates the resources needed for cuVS
@@ -43,9 +43,11 @@ public CuVSResourcesImpl(Path tempDirectory) {
     this.tempDirectory = tempDirectory;
     try (var localArena = Arena.ofConfined()) {
       var resourcesMemorySegment = localArena.allocate(cuvsResources_t);
-      int returnValue = cuvsResourcesCreate(resourcesMemorySegment);
-      checkCuVSError(returnValue, "cuvsResourcesCreate");
+      checkCuVSError(cuvsResourcesCreate(resourcesMemorySegment), "cuvsResourcesCreate");
       this.resourceHandle = resourcesMemorySegment.get(cuvsResources_t, 0);
+      var deviceIdPtr = localArena.allocate(C_INT);
+      checkCuVSError(cuvsDeviceIdGet(resourceHandle, deviceIdPtr), "cuvsDeviceIdGet");
+      this.deviceId = deviceIdPtr.get(C_INT, 0);
       this.access =
           new ScopedAccess() {
             @Override
@@ -64,6 +66,11 @@ public ScopedAccess access() {
     return this.access;
   }
 
+  @Override
+  public int deviceId() {
+    return this.deviceId;
+  }
+
   @Override
   public void close() {
     synchronized (this) {

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

@@ -0,0 +1,121 @@
+/*
+ * Copyright (c) 2025, NVIDIA CORPORATION.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package com.nvidia.cuvs.internal;
+
+import static com.nvidia.cuvs.internal.common.LinkerHelper.C_INT;
+import static com.nvidia.cuvs.internal.common.Util.checkCudaError;
+import static com.nvidia.cuvs.internal.panama.headers_h.cudaMemGetInfo;
+import static com.nvidia.cuvs.internal.panama.headers_h_1.*;
+
+import com.nvidia.cuvs.CuVSResources;
+import com.nvidia.cuvs.CuVSResourcesInfo;
+import com.nvidia.cuvs.GPUInfo;
+import com.nvidia.cuvs.GPUInfoProvider;
+import com.nvidia.cuvs.internal.panama.cudaDeviceProp;
+import java.lang.foreign.Arena;
+import java.lang.foreign.MemorySegment;
+import java.util.ArrayList;
+import java.util.List;
+
+public class GPUInfoProviderImpl implements GPUInfoProvider {
+
+  // Lazy initialization for list of available GPUs.
+  private static class AvailableGpuInitializer {
+
+    // Available GPUs are initialized only once when first accessed.
+    // This is assumed to be invariant for the lifetime of the program.
+    static final List<GPUInfo> AVAILABLE_GPUS = getAvailableGpusInfo();
+
+    private static List<GPUInfo> getAvailableGpusInfo() {
+      try (var localArena = Arena.ofConfined()) {
+
+        MemorySegment numGpus = localArena.allocate(C_INT);
+        int returnValue = cudaGetDeviceCount(numGpus);
+        checkCudaError(returnValue, "cudaGetDeviceCount");
+
+        int numGpuCount = numGpus.get(C_INT, 0);
+        List<GPUInfo> gpuInfoArr = new ArrayList<GPUInfo>();
+
+        MemorySegment deviceProp = cudaDeviceProp.allocate(localArena);
+
+        for (int i = 0; i < numGpuCount; i++) {
+          returnValue = cudaGetDeviceProperties_v2(deviceProp, i);
+          checkCudaError(returnValue, "cudaGetDeviceProperties_v2");
+
+          GPUInfo gpuInfo =
+              new GPUInfo(
+                  i,
+                  cudaDeviceProp.name(deviceProp).getString(0),
+                  cudaDeviceProp.totalGlobalMem(deviceProp),
+                  cudaDeviceProp.major(deviceProp),
+                  cudaDeviceProp.minor(deviceProp),
+                  cudaDeviceProp.asyncEngineCount(deviceProp) > 0,
+                  cudaDeviceProp.concurrentKernels(deviceProp) > 0);
+
+          gpuInfoArr.add(gpuInfo);
+        }
+        return gpuInfoArr;
+      }
+    }
+  }
+
+  private static boolean hasMinimumCapability(GPUInfo gpuInfo) {
+    return gpuInfo.computeCapabilityMajor() > GPUInfoProvider.MIN_COMPUTE_CAPABILITY_MAJOR
+        || (gpuInfo.computeCapabilityMajor() == GPUInfoProvider.MIN_COMPUTE_CAPABILITY_MAJOR
+            && gpuInfo.computeCapabilityMinor() >= GPUInfoProvider.MIN_COMPUTE_CAPABILITY_MINOR);
+  }
+
+  @Override
+  public List<GPUInfo> availableGPUs() {
+    return AvailableGpuInitializer.AVAILABLE_GPUS;
+  }
+
+  @Override
+  public List<GPUInfo> compatibleGPUs() {
+    List<GPUInfo> compatibleGPUs = new ArrayList<>();
+    long minDeviceMemoryInBytes = 1024L * 1024L * GPUInfoProvider.MIN_DEVICE_MEMORY_IN_MB;
+    for (GPUInfo gpuInfo : AvailableGpuInitializer.AVAILABLE_GPUS) {
+      if (hasMinimumCapability(gpuInfo)
+          && gpuInfo.totalDeviceMemoryInBytes() >= minDeviceMemoryInBytes) {
+        compatibleGPUs.add(gpuInfo);
+      }
+    }
+    return compatibleGPUs;
+  }
+
+  @Override
+  public CuVSResourcesInfo getCurrentInfo(CuVSResources resources) {
+    try (var localArena = Arena.ofConfined()) {
+      var deviceIdPtr = localArena.allocate(C_INT);
+      checkCudaError(cudaGetDevice(deviceIdPtr), "cudaGetDevice");
+      var currentDeviceId = deviceIdPtr.get(C_INT, 0);
+
+      if (resources.deviceId() != currentDeviceId) {
+        checkCudaError(cudaSetDevice(resources.deviceId()), "cudaSetDevice");
+      }
+
+      MemorySegment freeMemoryPtr = localArena.allocate(size_t);
+      MemorySegment totalMemoryPtr = localArena.allocate(size_t);
+      checkCudaError(cudaMemGetInfo(freeMemoryPtr, totalMemoryPtr), "cudaMemGetInfo");
+
+      if (resources.deviceId() != currentDeviceId) {
+        checkCudaError(cudaSetDevice(currentDeviceId), "cudaSetDevice");
+      }
+
+      return new CuVSResourcesInfo(freeMemoryPtr.get(size_t, 0), totalMemoryPtr.get(size_t, 0));
+    }
+  }
+}