Use the CUDA driver for cuda.core device enumeration (#2674)

* Use CUDA driver for CUDA device enumeration

* Adapt CUDA device enumeration to current main

Update the newer foreign-context test to use the CUDA-visible device count and regenerate the Device stub after applying the original public PR #2533 change.

---------

Co-authored-by: isvoid <isVoid@users.noreply.github.com>

Author: rwgk

cuda_core/cuda/core/_device.pyi

@@ -500,6 +500,8 @@ class Device:
         tuple of Device
             A tuple containing instances of available devices.
         """
+    @classmethod
+    def _get_all_devices_from_cuda_driver(cls): ...
     def to_system_device(self) -> cuda.core.system.Device:
         """
         Get the corresponding :class:`cuda.core.system.Device` (which is used

cuda_core/cuda/core/_device.pyx

@@ -1033,9 +1033,12 @@ class Device:
         tuple of Device
             A tuple containing instances of available devices.
         """
-        from cuda.core import system
-        total = system.get_num_devices()
-        return tuple(cls(device_id) for device_id in range(total))
+        return cls._get_all_devices_from_cuda_driver()
+
+    @classmethod
+    def _get_all_devices_from_cuda_driver(cls):
+        Device_ensure_cuda_initialized()
+        return tuple(Device_ensure_tls_devices(cls))
 
     def to_system_device(self) -> 'cuda.core.system.Device':
         """

cuda_core/docs/source/release/1.2.0-notes.rst

@@ -172,6 +172,11 @@ Fixes and enhancements
   not worth a warning.
   (closes `#2640 <https://github.com/NVIDIA/cuda-python/issues/2640>`__)
 
+- CUDA device enumeration now queries the CUDA driver rather than using the
+  NVML system-device count. This prevents non-CUDA accelerators, such as an NPU,
+  from being treated as CUDA devices by :meth:`Device.get_all_devices`, examples,
+  and tests.
+
 Deprecation Notices
 -------------------
 

cuda_core/examples/show_device_properties.py

@@ -15,7 +15,7 @@
 
 import sys
 
-from cuda.core import Device, system
+from cuda.core import Device
 
 
 # Convert boolean to YES or NO string
@@ -219,11 +219,12 @@ def print_device_properties(properties):
 
 # Print info about all CUDA devices in the system
 def show_device_properties():
-    ndev = system.get_num_devices()
+    devices = Device.get_all_devices()
+    ndev = len(devices)
     print(f"Number of GPUs: {ndev}")
 
-    for device_id in range(ndev):
-        device = Device(device_id)
+    for device in devices:
+        device_id = device.device_id
         print(f"DEVICE {device.name} (id={device_id})")
 
         device.set_current()

cuda_core/examples/simple_multi_gpu_example.py

@@ -17,7 +17,7 @@
 
 import cupy as cp
 
-from cuda.core import Device, LaunchConfig, Program, ProgramOptions, launch, system
+from cuda.core import Device, LaunchConfig, Program, ProgramOptions, launch
 
 dtype = cp.float32
 size = 50000
@@ -35,7 +35,7 @@ def __cuda_stream__(self):
 
 
 def main():
-    if system.get_num_devices() < 2:
+    if len(Device.get_all_devices()) < 2:
         print("this example requires at least 2 GPUs", file=sys.stderr)
         sys.exit(1)
 

cuda_core/tests/example_tests/test_basic_examples.py

@@ -13,7 +13,7 @@
 import pytest
 from cuda_python_test_helpers.pep723 import has_package_requirements_or_skip
 
-from cuda.core import Device, ManagedMemoryResource, system
+from cuda.core import Device, ManagedMemoryResource
 from cuda.core._program import _can_load_generated_ptx
 
 
@@ -22,7 +22,7 @@ def has_compute_capability_9_or_higher() -> bool:
 
 
 def has_multiple_devices() -> bool:
-    return system.get_num_devices() >= 2
+    return len(Device.get_all_devices()) >= 2
 
 
 def has_display() -> bool:

cuda_core/tests/system/test_system_device.py

@@ -15,6 +15,7 @@
 import helpers
 import pytest
 
+from cuda.core import Device as CudaDevice
 from cuda.core import system
 from cuda.core.system import typing
 
@@ -128,7 +129,8 @@ def test_numa_node_id(subtests):
 
 
 def test_device_cuda_compute_capability():
-    for device in system.Device.get_all_devices():
+    for cuda_device in CudaDevice.get_all_devices():
+        device = cuda_device.to_system_device()
         cuda_compute_capability = device.cuda_compute_capability
         assert isinstance(cuda_compute_capability, tuple)
         assert len(cuda_compute_capability) == 2
@@ -309,7 +311,8 @@ def test_device_brand():
 
 
 def test_device_pci_bus_id():
-    for device in system.Device.get_all_devices():
+    for cuda_device in CudaDevice.get_all_devices():
+        device = cuda_device.to_system_device()
         pci_bus_id = device.pci_info.bus_id
         assert isinstance(pci_bus_id, str)
 
@@ -852,7 +855,8 @@ def test_pstates(subtests):
 
 
 def test_compute_running_processes(subtests):
-    for device in system.Device.get_all_devices():
+    for cuda_device in CudaDevice.get_all_devices():
+        device = cuda_device.to_system_device()
         with subtests.test(device_index=device.index):
             with unsupported_before(device, "FERMI"):
                 processes = device.compute_running_processes

cuda_core/tests/system/test_system_system.py

@@ -9,6 +9,7 @@
 from cuda_python_test_helpers.arch_check import skip_if_nvml_unsupported
 
 from cuda.bindings import driver
+from cuda.core import Device as CudaDevice
 from cuda.core import system
 from cuda.core._utils.cuda_utils import handle_return
 
@@ -57,8 +58,9 @@ def test_nvml_version():
 
 @skip_if_nvml_unsupported
 def test_get_process_name():
-    for device in system.Device.get_all_devices():
-        x = device.compute_running_processes
+    for cuda_device in CudaDevice.get_all_devices():
+        device = cuda_device.to_system_device()
+        _ = device.compute_running_processes
 
     try:
         process_name = system.get_process_name(os.getpid())

cuda_core/tests/test_green_context.py

@@ -291,12 +291,12 @@ def test_configure_scope_with_enum(self, wq_resource, scope):
         assert wq_resource.sharing_scope is scope
 
     def test_device_id_matches_source_multi_gpu(self):
-        from cuda.core import Device, system
+        from cuda.core import Device
 
-        if system.get_num_devices() < 2:
+        devices = Device.get_all_devices()
+        if len(devices) < 2:
             pytest.skip("requires 2+ GPUs")
-        dev0 = Device(0)
-        dev1 = Device(1)
+        dev0, dev1 = devices[:2]
         try:
             wq0 = dev0.resources.workqueue
             wq1 = dev1.resources.workqueue

cuda_core/tests/test_memory.py

@@ -45,9 +45,6 @@
     VirtualMemoryResource,
     VirtualMemoryResourceOptions,
 )
-from cuda.core import (
-    system as ccx_system,
-)
 from cuda.core._dlpack import DLDeviceType
 from cuda.core._memory._ipc import IPCBufferDescriptor
 from cuda.core._stream import default_stream
@@ -402,7 +399,7 @@ def test_buffer_external_host():
 
 @pytest.mark.parametrize("change_device", [True, False])
 def test_buffer_external_device(change_device):
-    n = ccx_system.get_num_devices()
+    n = len(Device.get_all_devices())
     if n < 1:
         pytest.skip("No devices found")
     dev_id = n - 1
@@ -426,7 +423,7 @@ def test_buffer_external_device(change_device):
 
 @pytest.mark.parametrize("change_device", [True, False])
 def test_buffer_external_pinned_alloc(change_device):
-    n = ccx_system.get_num_devices()
+    n = len(Device.get_all_devices())
     if n < 1:
         pytest.skip("No devices found")
     dev_id = n - 1
@@ -451,7 +448,7 @@ def test_buffer_external_pinned_alloc(change_device):
 
 @pytest.mark.parametrize("change_device", [True, False])
 def test_buffer_external_pinned_registered(change_device):
-    n = ccx_system.get_num_devices()
+    n = len(Device.get_all_devices())
     if n < 1:
         pytest.skip("No devices found")
     dev_id = n - 1
@@ -484,7 +481,7 @@ def test_buffer_external_pinned_registered(change_device):
 
 @pytest.mark.parametrize("change_device", [True, False])
 def test_buffer_external_managed(change_device):
-    n = ccx_system.get_num_devices()
+    n = len(Device.get_all_devices())
     if n < 1:
         pytest.skip("No devices found")
     dev_id = n - 1
@@ -760,7 +757,7 @@ def test_mr_deallocation_without_current_context(init_cuda, capsys, replace_stre
 @pytest.mark.parametrize("replace_stream", [False, True])
 def test_mr_deallocation_with_foreign_context(capsys, replace_stream):
     """MR-backed Buffer teardown switches away from an unrelated current context."""
-    if ccx_system.get_num_devices() < 2:
+    if len(Device.get_all_devices()) < 2:
         pytest.skip("Test requires at least 2 GPUs")
 
     alloc_dev = Device(0)