Switch to use CUDA driver APIs in Device constructor

cuda_core/cuda/core/experimental/_device.py

@@ -13,6 +13,7 @@
 from cuda.core.experimental._utils.cuda_utils import (
  ComputeCapability,
  CUDAError,
+ _check_driver_error,
  driver,
  handle_return,
  precondition,
@@ -930,6 +931,10 @@ def multicast_supported(self) -> bool:
  return bool(self._get_cached_attribute(driver.CUdevice_attribute.CU_DEVICE_ATTRIBUTE_MULTICAST_SUPPORTED))
 
 
+_SUCCESS = driver.CUresult.CUDA_SUCCESS
+_INVALID_CTX = driver.CUresult.CUDA_ERROR_INVALID_CONTEXT
+
+
 class Device:
  """Represent a GPU and act as an entry point for cuda.core features.
 
@@ -959,7 +964,7 @@ class Device:
 
  __slots__ = ("_id", "_mr", "_has_inited", "_properties")
 
- def __new__(cls, device_id=None):
+ def __new__(cls, device_id: Optional[int] = None):
  global _is_cuInit
  if _is_cuInit is False:
  with _lock:
@@ -968,26 +973,34 @@ def __new__(cls, device_id=None):
 
  # important: creating a Device instance does not initialize the GPU!
  if device_id is None:
- device_id = handle_return(runtime.cudaGetDevice())
- assert_type(device_id, int)
- else:
- total = handle_return(runtime.cudaGetDeviceCount())
- assert_type(device_id, int)
- if not (0 <= device_id < total):
- raise ValueError(f"device_id must be within [0, {total}), got {device_id}")
+ err, dev = driver.cuCtxGetDevice()
+ if err == _SUCCESS:
+ device_id = int(dev)
+ elif err == _INVALID_CTX:
+ ctx = handle_return(driver.cuCtxGetCurrent())
+ assert int(ctx) == 0
+ device_id = 0 # cudart behavior
+ else:
+ _check_driver_error(err)
+ elif device_id < 0:
+ raise ValueError(f"device_id must be >= 0, got {device_id}")
 
  # ensure Device is singleton
- if not hasattr(_tls, "devices"):
- total = handle_return(runtime.cudaGetDeviceCount())
- _tls.devices = []
+ try:
+ devices = _tls.devices
+ except AttributeError:
+ total = handle_return(driver.cuDeviceGetCount())
+ devices = _tls.devices = []
  for dev_id in range(total):
  dev = super().__new__(cls)
  dev._id = dev_id
  # If the device is in TCC mode, or does not support memory pools for some other reason,
  # use the SynchronousMemoryResource which does not use memory pools.
  if (
  handle_return(
- runtime.cudaDeviceGetAttribute(runtime.cudaDeviceAttr.cudaDevAttrMemoryPoolsSupported, 0)
+ driver.cuDeviceGetAttribute(
+ driver.CUdevice_attribute.CU_DEVICE_ATTRIBUTE_MEMORY_POOLS_SUPPORTED, dev_id
+ )
  )
  ) == 1:
  dev._mr = _DefaultAsyncMempool(dev_id)
@@ -996,9 +1009,12 @@ def __new__(cls, device_id=None):
 
  dev._has_inited = False
  dev._properties = None
-_tls.devices.append(dev)
+ devices.append(dev)
 
- return _tls.devices[device_id]
+ try:
+ return devices[device_id]
+ except IndexError:
+ raise ValueError(f"device_id must be within [0, {len(devices)}), got {device_id}") from None
 
  def _check_context_initialized(self, *args, **kwargs):
  if not self._has_inited:

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

@@ -22,7 +22,7 @@ New features
 
 - :class:`Kernel` adds :attr:`Kernel.num_arguments` and :attr:`Kernel.arguments_info` for introspection of kernel arguments. (#612)
 - Add pythonic access to kernel occupancy calculation functions via :attr:`Kernel.occupancy`. (#648)
-- Support launching cooperative kernels by setting :property:`LaunchConfig.cooperative_launch` to `True`.
+- Support launching cooperative kernels by setting :attr:`LaunchConfig.cooperative_launch` to `True`.
 - A name can be assigned to :class:`ObjectCode` instances generated by both :class:`Program` and :class:`Linker` through their respective
  options.
 
@@ -34,5 +34,6 @@ New examples
 Fixes and enhancements
 ----------------------
 
-- An :class:`Event` can now be used to look up its corresponding device and context using the ``.device`` and ``.context`` attributes respectively.
+- Look-up of the :attr:`Event.device` and :attr:`Event.context` (the device and CUDA context where an event was created from) is now possible.
 - The :func:`launch` function's handling of fp16 scalars was incorrect and is fixed.
+- The :class:`Device` constructor is made faster.