Merge branch 'main' into IPC

Author: keenan-simpson

cuda_core/cuda/core/experimental/_context.py

@@ -4,7 +4,8 @@
 
 from dataclasses import dataclass
 
-from cuda.core.experimental._utils import driver
+from cuda.core.experimental._utils.clear_error_support import assert_type
+from cuda.core.experimental._utils.cuda_utils import driver
 
 
 @dataclass
@@ -20,7 +21,7 @@ def __new__(self, *args, **kwargs):
 
     @classmethod
     def _from_ctx(cls, obj, dev_id):
-        assert isinstance(obj, driver.CUcontext)
+        assert_type(obj, driver.CUcontext)
         ctx = super().__new__(cls)
         ctx._handle = obj
         ctx._id = dev_id

cuda_core/cuda/core/experimental/_device.py

@@ -3,12 +3,22 @@
 # SPDX-License-Identifier: LicenseRef-NVIDIA-SOFTWARE-LICENSE
 
 import threading
-from typing import Union
+from typing import Optional, Union
 
 from cuda.core.experimental._context import Context, ContextOptions
 from cuda.core.experimental._memory import AsyncMempool, Buffer, MemoryResource, _SynchronousMemoryResource
+from cuda.core.experimental._event import Event, EventOptions
+
 from cuda.core.experimental._stream import Stream, StreamOptions, default_stream
-from cuda.core.experimental._utils import ComputeCapability, CUDAError, driver, handle_return, precondition, runtime
+from cuda.core.experimental._utils.clear_error_support import assert_type
+from cuda.core.experimental._utils.cuda_utils import (
+    ComputeCapability,
+    CUDAError,
+    driver,
+    handle_return,
+    precondition,
+    runtime,
+)
 
 _tls = threading.local()
 _lock = threading.Lock()
@@ -949,10 +959,11 @@ 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 isinstance(device_id, int), f"{device_id=}"
+            assert_type(device_id, int)
         else:
             total = handle_return(runtime.cudaGetDeviceCount())
-            if not isinstance(device_id, int) or not (0 <= device_id < total):
+            assert_type(device_id, int)
+            if not (0 <= device_id < total):
                 raise ValueError(f"device_id must be within [0, {total}), got {device_id}")
 
         # ensure Device is singleton
@@ -981,7 +992,9 @@ def __new__(cls, device_id=None):
 
     def _check_context_initialized(self, *args, **kwargs):
         if not self._has_inited:
-            raise CUDAError("the device is not yet initialized, perhaps you forgot to call .set_current() first?")
+            raise CUDAError(
+                f"Device {self._id} is not yet initialized, perhaps you forgot to call .set_current() first?"
+            )
 
     @property
     def device_id(self) -> int:
@@ -1053,7 +1066,8 @@ def context(self) -> Context:
 
         """
         ctx = handle_return(driver.cuCtxGetCurrent())
-        assert int(ctx) != 0
+        if int(ctx) == 0:
+            raise CUDAError("No context is bound to the calling CPU thread.")
         return Context._from_ctx(ctx, self._id)
 
     @property
@@ -1063,8 +1077,7 @@ def memory_resource(self) -> MemoryResource:
 
     @memory_resource.setter
     def memory_resource(self, mr):
-        if not isinstance(mr, MemoryResource):
-            raise TypeError
+        assert_type(mr, MemoryResource)
         self._mr = mr
 
     @property
@@ -1118,12 +1131,11 @@ def set_current(self, ctx: Context = None) -> Union[Context, None]:
 
         """
         if ctx is not None:
-            if not isinstance(ctx, Context):
-                raise TypeError("a Context object is required")
+            assert_type(ctx, Context)
             if ctx._id != self._id:
                 raise RuntimeError(
-                    "the provided context was created on a different "
-                    f"device {ctx._id} other than the target {self._id}"
+                    "the provided context was created on the device with"
+                    f" id={ctx._id}, which is different from the target id={self._id}"
                 )
             prev_ctx = handle_return(driver.cuCtxPopCurrent())
             handle_return(driver.cuCtxPushCurrent(ctx._handle))
@@ -1165,7 +1177,7 @@ def create_context(self, options: ContextOptions = None) -> Context:
             Newly created context object.
 
         """
-        raise NotImplementedError("TODO")
+        raise NotImplementedError("WIP: https://github.com/NVIDIA/cuda-python/issues/189")
 
     @precondition(_check_context_initialized)
     def create_stream(self, obj=None, options: StreamOptions = None) -> Stream:
@@ -1198,6 +1210,27 @@ def create_stream(self, obj=None, options: StreamOptions = None) -> Stream:
         """
         return Stream._init(obj=obj, options=options)
 
+    @precondition(_check_context_initialized)
+    def create_event(self, options: Optional[EventOptions] = None) -> Event:
+        """Create an Event object without recording it to a Stream.
+
+        Note
+        ----
+        Device must be initialized.
+
+        Parameters
+        ----------
+        options : :obj:`EventOptions`, optional
+            Customizable dataclass for event creation options.
+
+        Returns
+        -------
+        :obj:`~_event.Event`
+            Newly created event object.
+
+        """
+        return Event._init(options)
+
     @precondition(_check_context_initialized)
     def allocate(self, size, stream=None) -> Buffer:
         """Allocate device memory from a specified stream.

cuda_core/cuda/core/experimental/_event.py

@@ -8,7 +8,7 @@
 from dataclasses import dataclass
 from typing import TYPE_CHECKING, Optional
 
-from cuda.core.experimental._utils import CUDAError, check_or_create_options, driver, handle_return
+from cuda.core.experimental._utils.cuda_utils import CUDAError, check_or_create_options, driver, handle_return
 
 if TYPE_CHECKING:
     import cuda.bindings
@@ -47,8 +47,21 @@ class Event:
     the last recorded stream.
 
     Events can be used to monitor device's progress, query completion
-    of work up to event's record, and help establish dependencies
-    between GPU work submissions.
+    of work up to event's record, help establish dependencies
+    between GPU work submissions, and record the elapsed time (in milliseconds)
+    on GPU:
+
+    .. code-block:: python
+
+        # To create events and record the timing:
+        s = Device().create_stream()
+        e1 = Device().create_event({"enable_timing": True})
+        e2 = Device().create_event({"enable_timing": True})
+        s.record(e1)
+        # ... run some GPU works ...
+        s.record(e2)
+        e2.sync()
+        print(f"time = {e2 - e1} milliseconds")
 
     Directly creating an :obj:`~_event.Event` is not supported due to ambiguity,
     and they should instead be created through a :obj:`~_stream.Stream` object.
@@ -88,14 +101,30 @@ def _init(cls, options: Optional[EventOptions] = None):
             flags |= driver.CUevent_flags.CU_EVENT_BLOCKING_SYNC
             self._busy_waited = True
         if options.support_ipc:
-            raise NotImplementedError("TODO")
+            raise NotImplementedError("WIP: https://github.com/NVIDIA/cuda-python/issues/103")
         self._mnff.handle = handle_return(driver.cuEventCreate(flags))
         return self
 
     def close(self):
         """Destroy the event."""
         self._mnff.close()
 
+    def __isub__(self, other):
+        return NotImplemented
+
+    def __rsub__(self, other):
+        return NotImplemented
+
+    def __sub__(self, other):
+        # return self - other (in milliseconds)
+        try:
+            timing = handle_return(driver.cuEventElapsedTime(other.handle, self.handle))
+        except CUDAError as e:
+            raise RuntimeError(
+                "Timing capability must be enabled in order to subtract two Events; timing is disabled by default."
+            ) from e
+        return timing
+
     @property
     def is_timing_disabled(self) -> bool:
         """Return True if the event does not record timing data, otherwise False."""
@@ -109,7 +138,7 @@ def is_sync_busy_waited(self) -> bool:
     @property
     def is_ipc_supported(self) -> bool:
         """Return True if this event can be used as an interprocess event, otherwise False."""
-        raise NotImplementedError("TODO")
+        raise NotImplementedError("WIP: https://github.com/NVIDIA/cuda-python/issues/103")
 
     def sync(self):
         """Synchronize until the event completes.
@@ -129,10 +158,9 @@ def is_done(self) -> bool:
         (result,) = driver.cuEventQuery(self._mnff.handle)
         if result == driver.CUresult.CUDA_SUCCESS:
             return True
-        elif result == driver.CUresult.CUDA_ERROR_NOT_READY:
+        if result == driver.CUresult.CUDA_ERROR_NOT_READY:
             return False
-        else:
-            raise CUDAError(f"unexpected error: {result}")
+        handle_return(result)
 
     @property
     def handle(self) -> cuda.bindings.driver.CUevent:

cuda_core/cuda/core/experimental/_launcher.py

@@ -9,7 +9,15 @@
 from cuda.core.experimental._kernel_arg_handler import ParamHolder
 from cuda.core.experimental._module import Kernel
 from cuda.core.experimental._stream import Stream
-from cuda.core.experimental._utils import CUDAError, check_or_create_options, driver, get_binding_version, handle_return
+from cuda.core.experimental._utils.clear_error_support import assert_type
+from cuda.core.experimental._utils.cuda_utils import (
+    CUDAError,
+    cast_to_3_tuple,
+    check_or_create_options,
+    driver,
+    get_binding_version,
+    handle_return,
+)
 
 # TODO: revisit this treatment for py313t builds
 _inited = False
@@ -59,41 +67,23 @@ class LaunchConfig:
 
     def __post_init__(self):
         _lazy_init()
-        self.grid = self._cast_to_3_tuple(self.grid)
-        self.block = self._cast_to_3_tuple(self.block)
+        self.grid = cast_to_3_tuple("LaunchConfig.grid", self.grid)
+        self.block = cast_to_3_tuple("LaunchConfig.block", self.block)
         # thread block clusters are supported starting H100
         if self.cluster is not None:
             if not _use_ex:
-                raise CUDAError("thread block clusters require cuda.bindings & driver 11.8+")
-            if Device().compute_capability < (9, 0):
-                raise CUDAError("thread block clusters are not supported on devices with compute capability < 9.0")
-            self.cluster = self._cast_to_3_tuple(self.cluster)
+                err, drvers = driver.cuDriverGetVersion()
+                drvers_fmt = f" (got driver version {drvers})" if err == driver.CUresult.CUDA_SUCCESS else ""
+                raise CUDAError(f"thread block clusters require cuda.bindings & driver 11.8+{drvers_fmt}")
+            cc = Device().compute_capability
+            if cc < (9, 0):
+                raise CUDAError(
+                    f"thread block clusters are not supported on devices with compute capability < 9.0 (got {cc})"
+                )
+            self.cluster = cast_to_3_tuple("LaunchConfig.cluster", self.cluster)
         if self.shmem_size is None:
             self.shmem_size = 0
 
-    def _cast_to_3_tuple(self, cfg):
-        if isinstance(cfg, int):
-            if cfg < 1:
-                raise ValueError
-            return (cfg, 1, 1)
-        elif isinstance(cfg, tuple):
-            size = len(cfg)
-            if size == 1:
-                cfg = cfg[0]
-                if cfg < 1:
-                    raise ValueError
-                return (cfg, 1, 1)
-            elif size == 2:
-                if cfg[0] < 1 or cfg[1] < 1:
-                    raise ValueError
-                return (*cfg, 1)
-            elif size == 3:
-                if cfg[0] < 1 or cfg[1] < 1 or cfg[2] < 1:
-                    raise ValueError
-                return cfg
-        else:
-            raise ValueError
-
 
 def launch(stream, config, kernel, *kernel_args):
     """Launches a :obj:`~_module.Kernel`
@@ -120,9 +110,10 @@ def launch(stream, config, kernel, *kernel_args):
         try:
             stream = Stream._init(stream)
         except Exception as e:
-            raise ValueError("stream must either be a Stream object or support __cuda_stream__") from e
-    if not isinstance(kernel, Kernel):
-        raise ValueError
+            raise ValueError(
+                f"stream must either be a Stream object or support __cuda_stream__ (got {type(stream)})"
+            ) from e
+    assert_type(kernel, Kernel)
     config = check_or_create_options(LaunchConfig, config, "launch config")
 
     # TODO: can we ensure kernel_args is valid/safe to use here?

cuda_core/cuda/core/experimental/_linker.py

@@ -16,7 +16,8 @@
 
 from cuda.core.experimental._device import Device
 from cuda.core.experimental._module import ObjectCode
-from cuda.core.experimental._utils import check_or_create_options, driver, handle_return, is_sequence
+from cuda.core.experimental._utils.clear_error_support import assert_type
+from cuda.core.experimental._utils.cuda_utils import check_or_create_options, driver, handle_return, is_sequence
 
 # TODO: revisit this treatment for py313t builds
 _driver = None  # populated if nvJitLink cannot be used
@@ -382,12 +383,12 @@ def __init__(self, *object_codes: ObjectCode, options: LinkerOptions = None):
         self._mnff = Linker._MembersNeededForFinalize(self, handle, use_nvjitlink)
 
         for code in object_codes:
-            assert isinstance(code, ObjectCode)
+            assert_type(code, ObjectCode)
             self._add_code_object(code)
 
     def _add_code_object(self, object_code: ObjectCode):
         data = object_code._module
-        assert isinstance(data, bytes)
+        assert_type(data, bytes)
         with _exception_manager(self):
             if _nvjitlink:
                 _nvjitlink.add_data(