Merge pull request #256 from ksimpson-work/add_system

Add the cuda.core.experimental.system singleton

Author: leofang

cuda_core/cuda/core/experimental/__init__.py

@@ -9,3 +9,8 @@
 from cuda.core.experimental._linker import Linker, LinkerOptions
 from cuda.core.experimental._program import Program
 from cuda.core.experimental._stream import Stream, StreamOptions
+from cuda.core.experimental._system import System
+
+system = System()
+__import__("sys").modules[__spec__.name + ".system"] = system
+del System

cuda_core/cuda/core/experimental/_linker.py

@@ -443,7 +443,7 @@ def link(self, target_type) -> ObjectCode:
         return ObjectCode(bytes(code), target_type)
 
     def get_error_log(self) -> str:
-        """ Get the error log generated by the linker.
+        """Get the error log generated by the linker.
 
         Returns
         -------

cuda_core/cuda/core/experimental/_system.py

@@ -0,0 +1,67 @@
+# Copyright (c) 2024, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED.
+#
+# SPDX-License-Identifier: LicenseRef-NVIDIA-SOFTWARE-LICENSE
+
+from typing import Tuple
+
+from cuda import cuda, cudart
+from cuda.core.experimental._device import Device
+from cuda.core.experimental._utils import handle_return
+
+
+class System:
+    """Provide information about the cuda system.
+    This class is a singleton and should not be instantiated directly.
+    """
+
+    _instance = None
+
+    def __new__(cls):
+        if cls._instance is None:
+            cls._instance = super().__new__(cls)
+        return cls._instance
+
+    def __init__(self):
+        if hasattr(self, "_initialized") and self._initialized:
+            return
+        self._initialized = True
+
+    @property
+    def driver_version(self) -> Tuple[int, int]:
+        """
+        Query the CUDA driver version.
+
+        Returns
+        -------
+        tuple of int
+            A 2-tuple of (major, minor) version numbers.
+        """
+        version = handle_return(cuda.cuDriverGetVersion())
+        major = version // 1000
+        minor = (version % 1000) // 10
+        return (major, minor)
+
+    @property
+    def num_devices(self) -> int:
+        """
+        Query the number of available GPUs.
+
+        Returns
+        -------
+        int
+            The number of available GPU devices.
+        """
+        return handle_return(cudart.cudaGetDeviceCount())
+
+    @property
+    def devices(self) -> tuple:
+        """
+        Query the available device instances.
+
+        Returns
+        -------
+        tuple of Device
+            A tuple containing instances of available devices.
+        """
+        total = self.num_devices
+        return tuple(Device(device_id) for device_id in range(total))

cuda_core/docs/source/api.rst

@@ -38,6 +38,17 @@ CUDA compilation toolchain
    LinkerOptions
 
 
+CUDA system information
+-----------------------
+
+.. autodata:: cuda.core.experimental.system.driver_version
+   :no-value:
+.. autodata:: cuda.core.experimental.system.num_devices
+   :no-value:
+.. autodata:: cuda.core.experimental.system.devices
+   :no-value:
+
+
 .. module:: cuda.core.experimental.utils
 
 Utility functions

cuda_core/docs/source/conf.py

@@ -91,3 +91,31 @@
 
 napoleon_google_docstring = False
 napoleon_numpy_docstring = True
+
+
+section_titles = ["Returns"]
+def autodoc_process_docstring(app, what, name, obj, options, lines):
+    if name.startswith("cuda.core.experimental.system"):
+        # patch the docstring (in lines) *in-place*. Should docstrings include section titles other than "Returns", 
+        # this will need to be modified to handle them.
+        attr = name.split(".")[-1]
+        from cuda.core.experimental._system import System
+
+        lines_new = getattr(System, attr).__doc__.split("\n")
+        formatted_lines = []
+        for line in lines_new:
+            title = line.strip()
+            if title in section_titles:
+                formatted_lines.append(line.replace(title, f".. rubric:: {title}"))
+            elif line.strip() == "-" * len(title):
+                formatted_lines.append(" " * len(title))
+            else:
+                formatted_lines.append(line)
+        n_pops = len(lines)
+        lines.extend(formatted_lines)
+        for _ in range(n_pops):
+            lines.pop(0)
+
+
+def setup(app):
+    app.connect("autodoc-process-docstring", autodoc_process_docstring)

cuda_core/docs/source/release.md

@@ -7,4 +7,5 @@ maxdepth: 3
 
     0.1.1 <release/0.1.1-notes>
     0.1.0 <release/0.1.0-notes>
+
 ```

cuda_core/docs/source/release/0.1.1-notes.md

@@ -9,6 +9,7 @@ Released on Dec XX, 2024
 - Add `Linker` that can link one or multiple `ObjectCode` instances generated by `Program`s. Under
   the hood, it uses either the nvJitLink or cuLink APIs depending on the CUDA version detected
   in the current environment.
+- Add a `cuda.core.experimental.system` module for querying system- or process- wide information.
 - Support TCC devices with a default synchronous memory resource to avoid the use of memory pools
 
 

cuda_core/tests/test_system.py

@@ -0,0 +1,37 @@
+try:
+    from cuda.bindings import driver, runtime
+except ImportError:
+    from cuda import cuda as driver
+    from cuda import cudart as runtime
+
+from cuda.core.experimental import Device, system
+from cuda.core.experimental._utils import handle_return
+
+
+def test_system_singleton():
+    system1 = system
+    system2 = system
+    assert id(system1) == id(system2), "system is not a singleton"
+
+
+def test_driver_version():
+    driver_version = system.driver_version
+    print(driver_version)
+    version = handle_return(driver.cuDriverGetVersion())
+    expected_driver_version = (version // 1000, (version % 1000) // 10)
+    assert driver_version == expected_driver_version, "Driver version does not match expected value"
+
+
+def test_num_devices():
+    num_devices = system.num_devices
+    expected_num_devices = handle_return(runtime.cudaGetDeviceCount())
+    assert num_devices == expected_num_devices, "Number of devices does not match expected value"
+
+
+def test_devices():
+    devices = system.devices
+    expected_num_devices = handle_return(runtime.cudaGetDeviceCount())
+    expected_devices = tuple(Device(device_id) for device_id in range(expected_num_devices))
+    assert len(devices) == len(expected_devices), "Number of devices does not match expected value"
+    for device, expected_device in zip(devices, expected_devices):
+        assert device.device_id == expected_device.device_id, "Device ID does not match expected value"