enchanced typing in buffer, MR, Stream, ...

Author: leofang

cuda_core/cuda/core/experimental/_device.py

@@ -8,7 +8,7 @@
 from cuda.core.experimental._context import Context, ContextOptions
 from cuda.core.experimental._event import Event, EventOptions
 from cuda.core.experimental._memory import Buffer, MemoryResource, _DefaultAsyncMempool, _SynchronousMemoryResource
-from cuda.core.experimental._stream import Stream, StreamOptions, default_stream
+from cuda.core.experimental._stream import IsStreamT, Stream, StreamOptions, default_stream
 from cuda.core.experimental._utils.clear_error_support import assert_type
 from cuda.core.experimental._utils.cuda_utils import (
     ComputeCapability,
@@ -1206,7 +1206,7 @@ def create_context(self, options: ContextOptions = None) -> Context:
         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:
+    def create_stream(self, obj: Optional[IsStreamT] = None, options: StreamOptions = None) -> Stream:
         """Create a Stream object.
 
         New stream objects can be created in two different ways:
@@ -1258,7 +1258,7 @@ def create_event(self, options: Optional[EventOptions] = None) -> Event:
         return Event._init(self._id, self.context._handle, options)
 
     @precondition(_check_context_initialized)
-    def allocate(self, size, stream=None) -> Buffer:
+    def allocate(self, size, stream: Optional[Stream] = None) -> Buffer:
         """Allocate device memory from a specified stream.
 
         Allocates device memory of `size` bytes on the specified `stream`

cuda_core/cuda/core/experimental/_launch_config.py

@@ -50,7 +50,8 @@ class LaunchConfig:
     shmem_size : int, optional
         Dynamic shared-memory size per thread block in bytes.
         (Default to size 0)
-
+    cooperative_launch : bool, optional
+        Whether this config can be used to launch a cooperative kernel.
     """
 
     # TODO: expand LaunchConfig to include other attributes

cuda_core/cuda/core/experimental/_launcher.py

@@ -2,11 +2,12 @@
 #
 # SPDX-License-Identifier: Apache-2.0
 
+from typing import Union
 
 from cuda.core.experimental._kernel_arg_handler import ParamHolder
 from cuda.core.experimental._launch_config import LaunchConfig, _to_native_launch_config
 from cuda.core.experimental._module import Kernel
-from cuda.core.experimental._stream import Stream
+from cuda.core.experimental._stream import IsStreamT, Stream, _try_to_get_stream_ptr
 from cuda.core.experimental._utils.clear_error_support import assert_type
 from cuda.core.experimental._utils.cuda_utils import (
     _reduce_3_tuple,
@@ -34,7 +35,7 @@ def _lazy_init():
     _inited = True
 
 
-def launch(stream, config, kernel, *kernel_args):
+def launch(stream: Union[Stream, IsStreamT], config: LaunchConfig, kernel: Kernel, *kernel_args):
     """Launches a :obj:`~_module.Kernel`
     object with launch-time configuration.
 
@@ -43,7 +44,7 @@ def launch(stream, config, kernel, *kernel_args):
     stream : :obj:`~_stream.Stream`
         The stream establishing the stream ordering semantic of a
         launch.
-    config : :obj:`~_launcher.LaunchConfig`
+    config : :obj:`LaunchConfig`
         Launch configurations inline with options provided by
         :obj:`~_launcher.LaunchConfig` dataclass.
     kernel : :obj:`~_module.Kernel`
@@ -55,13 +56,15 @@ def launch(stream, config, kernel, *kernel_args):
     """
     if stream is None:
         raise ValueError("stream cannot be None, stream must either be a Stream object or support __cuda_stream__")
-    if not isinstance(stream, Stream):
+    try:
+        stream_handle = stream.handle
+    except AttributeError:
         try:
-            stream = Stream._init(stream)
-        except Exception as e:
+            stream_handle = _try_to_get_stream_ptr(stream)
+        except Exception:
             raise ValueError(
                 f"stream must either be a Stream object or support __cuda_stream__ (got {type(stream)})"
-            ) from e
+            ) from None
     assert_type(kernel, Kernel)
     _lazy_init()
     config = check_or_create_options(LaunchConfig, config, "launch config")
@@ -78,15 +81,15 @@ def launch(stream, config, kernel, *kernel_args):
     # rich.
     if _use_ex:
         drv_cfg = _to_native_launch_config(config)
-        drv_cfg.hStream = stream.handle
+        drv_cfg.hStream = stream_handle
         if config.cooperative_launch:
             _check_cooperative_launch(kernel, config, stream)
         handle_return(driver.cuLaunchKernelEx(drv_cfg, int(kernel._handle), args_ptr, 0))
     else:
         # TODO: check if config has any unsupported attrs
         handle_return(
             driver.cuLaunchKernel(
-                int(kernel._handle), *config.grid, *config.block, config.shmem_size, stream.handle, args_ptr, 0
+                int(kernel._handle), *config.grid, *config.block, config.shmem_size, stream_handle, args_ptr, 0
             )
         )
 

cuda_core/cuda/core/experimental/_memory.py

@@ -9,17 +9,19 @@
 from typing import Optional, Tuple, TypeVar, Union
 
 from cuda.core.experimental._dlpack import DLDeviceType, make_py_capsule
-from cuda.core.experimental._stream import default_stream
+from cuda.core.experimental._stream import Stream, default_stream
 from cuda.core.experimental._utils.cuda_utils import driver, handle_return
 
-PyCapsule = TypeVar("PyCapsule")
-
 
 # TODO: define a memory property mixin class and make Buffer and
 # MemoryResource both inherit from it
 
+
+PyCapsule = TypeVar("PyCapsule")
+"""Represent the capsule type."""
+
 DevicePointerT = Union[driver.CUdeviceptr, int, None]
-"""A type union of `Cudeviceptr`, `int` and `None` for hinting Buffer.handle."""
+"""A type union of :obj:`~driver.CUdeviceptr`, `int` and `None` for hinting :attr:`Buffer.handle`."""
 
 
 class Buffer:
@@ -29,19 +31,7 @@ class Buffer:
     different memory resources are to give access to their memory
     allocations.
 
-    Support for data interchange mechanisms are provided by
-    establishing both the DLPack and the Python-level buffer
-    protocols.
-
-    Parameters
-    ----------
-    ptr : Any
-        Allocated buffer handle object
-    size : Any
-        Memory size of the buffer
-    mr : :obj:`~_memory.MemoryResource`, optional
-        Memory resource associated with the buffer
-
+    Support for data interchange mechanisms are provided by DLPack.
     """
 
     class _MembersNeededForFinalize:
@@ -64,22 +54,26 @@ def close(self, stream=None):
     # TODO: handle ownership? (_mr could be None)
     __slots__ = ("__weakref__", "_mnff")
 
-    def __init__(self, ptr, size, mr: MemoryResource = None):
+    def __new__(self, *args, **kwargs):
+        raise RuntimeError("Buffer objects cannot be instantiated directly. Please use MemoryResource APIs.")
+
+    @classmethod
+    def _init(cls, ptr: DevicePointerT, size: int, mr: Optional[MemoryResource] = None):
+        self = super().__new__(cls)
         self._mnff = Buffer._MembersNeededForFinalize(self, ptr, size, mr)
+        return self
 
-    def close(self, stream=None):
+    def close(self, stream: Stream = None):
         """Deallocate this buffer asynchronously on the given stream.
 
         This buffer is released back to their memory resource
         asynchronously on the given stream.
 
         Parameters
         ----------
-        stream : Any, optional
-            The stream object with a __cuda_stream__ protocol to
-            use for asynchronous deallocation. Defaults to using
-            the default stream.
-
+        stream : Stream, optional
+            The stream object to use for asynchronous deallocation. If not set,
+            the current default is to the default stream.
         """
         self._mnff.close(stream)
 
@@ -95,7 +89,7 @@ def handle(self) -> DevicePointerT:
         return self._mnff.ptr
 
     @property
-    def size(self):
+    def size(self) -> int:
         """Return the memory size of this buffer."""
         return self._mnff.size
 
@@ -125,7 +119,7 @@ def device_id(self) -> int:
             return self._mnff.mr.device_id
         raise NotImplementedError("WIP: Currently this property only supports buffers with associated MemoryResource")
 
-    def copy_to(self, dst: Buffer = None, *, stream) -> Buffer:
+    def copy_to(self, dst: Buffer = None, *, stream: Stream) -> Buffer:
         """Copy from this buffer to the dst buffer asynchronously on the given stream.
 
         Copies the data from this buffer to the provided dst buffer.
@@ -136,7 +130,7 @@ def copy_to(self, dst: Buffer = None, *, stream) -> Buffer:
         ----------
         dst : :obj:`~_memory.Buffer`
             Source buffer to copy data from
-        stream : Any
+        stream : Stream
             Keyword argument specifying the stream for the
             asynchronous copy
 
@@ -154,14 +148,14 @@ def copy_to(self, dst: Buffer = None, *, stream) -> Buffer:
         handle_return(driver.cuMemcpyAsync(dst._mnff.ptr, self._mnff.ptr, self._mnff.size, stream.handle))
         return dst
 
-    def copy_from(self, src: Buffer, *, stream):
+    def copy_from(self, src: Buffer, *, stream: Stream):
         """Copy from the src buffer to this buffer asynchronously on the given stream.
 
         Parameters
         ----------
         src : :obj:`~_memory.Buffer`
             Source buffer to copy data from
-        stream : Any
+        stream : Stream
             Keyword argument specifying the stream for the
             asynchronous copy
 
@@ -219,55 +213,117 @@ def __release_buffer__(self, buffer: memoryview, /):
         # Supporting method paired with __buffer__.
         raise NotImplementedError("WIP: Buffer.__release_buffer__ hasn't been implemented yet.")
 
+    @staticmethod
+    def from_handle(ptr: DevicePointerT, size: int, mr: Optional[MemoryResource] = None) -> Buffer:
+        """Create a new :class:`Buffer` object from a pointer.
+
+        Parameters
+        ----------
+        ptr : :obj:`~_memory.DevicePointerT`
+            Allocated buffer handle object
+        size : int
+            Memory size of the buffer
+        mr : :obj:`~_memory.MemoryResource`, optional
+            Memory resource associated with the buffer
+        """
+        return Buffer._init(ptr, size, mr=mr)
+
 
 class MemoryResource(abc.ABC):
+    """Abstract base class for memory resources that manage allocation and deallocation of buffers.
+
+    Subclasses must implement methods for allocating and deallocation, as well as properties
+    associated with this memory resource from which all allocated buffers will inherit. (Since
+    all :class:`Buffer` instances allocated and returned by the :meth:`allocate` method would
+    hold a reference to self, the buffer properties are retrieved simply by looking up the underlying
+    memory resource's respective property.)
+    """
+
     __slots__ = ("_handle",)
 
     @abc.abstractmethod
-    def __init__(self, *args, **kwargs): ...
+    def __init__(self, *args, **kwargs):
+        """Initialize the memory resource.
+
+        Subclasses may use additional arguments to configure the resource.
+        """
+        ...
 
     @abc.abstractmethod
-    def allocate(self, size, stream=None) -> Buffer: ...
+    def allocate(self, size: int, stream: Stream = None) -> Buffer:
+        """Allocate a buffer of the requested size.
+
+        Parameters
+        ----------
+        size : int
+            The size of the buffer to allocate, in bytes.
+        stream : object, optional
+            The stream on which to perform the allocation asynchronously.
+            If None, allocation is synchronous.
+
+        Returns
+        -------
+        Buffer
+            The allocated buffer object, which can be used for device or host operations
+            depending on the resource's properties.
+        """
+        ...
 
     @abc.abstractmethod
-    def deallocate(self, ptr, size, stream=None): ...
+    def deallocate(self, ptr: DevicePointerT, size: int, stream: Stream = None):
+        """Deallocate a buffer previously allocated by this resource.
+
+        Parameters
+        ----------
+        ptr : object
+            The pointer or handle to the buffer to deallocate.
+        size : int
+            The size of the buffer to deallocate, in bytes.
+        stream : object, optional
+            The stream on which to perform the deallocation asynchronously.
+            If None, deallocation is synchronous.
+        """
+        ...
 
     @property
     @abc.abstractmethod
     def is_device_accessible(self) -> bool:
-        # Check if the buffers allocated from this MR can be accessed from
-        # GPUs.
+        """bool: True if buffers allocated by this resource can be accessed on the device."""
         ...
 
     @property
     @abc.abstractmethod
     def is_host_accessible(self) -> bool:
-        # Check if the buffers allocated from this MR can be accessed from
-        # CPUs.
+        """bool: True if buffers allocated by this resource can be accessed on the host."""
         ...
 
     @property
     @abc.abstractmethod
     def device_id(self) -> int:
-        # Return the device ID if this MR is for single devices. Raise an
-        # exception if it is not.
+        """int: The device ordinal for which this memory resource is responsible.
+
+        Raises
+        ------
+        RuntimeError
+            If the resource is not bound to a specific device.
+        """
         ...
 
 
 class _DefaultAsyncMempool(MemoryResource):
     __slots__ = ("_dev_id",)
 
-    def __init__(self, dev_id):
+    def __init__(self, dev_id: int):
         self._handle = handle_return(driver.cuDeviceGetMemPool(dev_id))
         self._dev_id = dev_id
 
-    def allocate(self, size, stream=None) -> Buffer:
+    def allocate(self, size: int, stream: Stream = None) -> Buffer:
         if stream is None:
             stream = default_stream()
         ptr = handle_return(driver.cuMemAllocFromPoolAsync(size, self._handle, stream.handle))
-        return Buffer(ptr, size, self)
+        return Buffer._init(ptr, size, self)
 
-    def deallocate(self, ptr, size, stream=None):
+    def deallocate(self, ptr: DevicePointerT, size: int, stream: Stream = None):
         if stream is None:
             stream = default_stream()
         handle_return(driver.cuMemFreeAsync(ptr, stream.handle))
@@ -290,11 +346,11 @@ def __init__(self):
         # TODO: support flags from cuMemHostAlloc?
         self._handle = None
 
-    def allocate(self, size, stream=None) -> Buffer:
+    def allocate(self, size: int, stream: Stream = None) -> Buffer:
         ptr = handle_return(driver.cuMemAllocHost(size))
-        return Buffer(ptr, size, self)
+        return Buffer._init(ptr, size, self)
 
-    def deallocate(self, ptr, size, stream=None):
+    def deallocate(self, ptr: DevicePointerT, size: int, stream: Stream = None):
         handle_return(driver.cuMemFreeHost(ptr))
 
     @property
@@ -319,7 +375,7 @@ def __init__(self, dev_id):
 
     def allocate(self, size, stream=None) -> Buffer:
         ptr = handle_return(driver.cuMemAlloc(size))
-        return Buffer(ptr, size, self)
+        return Buffer._init(ptr, size, self)
 
     def deallocate(self, ptr, size, stream=None):
         if stream is None: