Polish green context API: docs, error handling, simplification

leofang · claude · leofang · commit eebd4cf5092d · 2026-04-27T13:42:00.000Z
- dev.create_context raises ValueError (not NotImplementedError) when
  options or resources are missing.
- Cache version checks (_check_green_ctx_support, _check_workqueue_support)
  at module level; raise ValueError instead of NotImplementedError.
- Simplify _device_resources.pyx: merge _as_uint and _count_to_sm_count
  into _to_sm_count; inline unsigned int casts for coscheduled params.
- Add green context classes to api.rst (Context, ContextOptions,
  DeviceResources, SMResource, SMResourceOptions, WorkqueueResource,
  WorkqueueResourceOptions).
- Update all docstrings to NumPy style with Attributes/Parameters/Returns
  sections matching the existing codebase convention.

Co-Authored-By: Claude Opus 4.6 (1M context) &lt;noreply@anthropic.com&gt;
diff --git a/cuda_core/cuda/core/_device.pyx b/cuda_core/cuda/core/_device.pyx
@@ -1302,11 +1302,15 @@ class Device:
         cdef GreenCtxHandle h_green
 
         if options is None:
-            raise NotImplementedError("WIP: https://github.com/NVIDIA/cuda-python/issues/189")
+            raise ValueError(
+                "options with device resources must be provided to create a green context"
+            )
 
         options = check_or_create_options(ContextOptions, options, "Context options")
         if options.resources is None:
-            raise NotImplementedError("WIP: https://github.com/NVIDIA/cuda-python/issues/189")
+            raise ValueError(
+                "ContextOptions.resources must be provided to create a green context"
+            )
 
         resources = tuple(options.resources)
         if len(resources) == 0:
@@ -1334,9 +1338,9 @@ class Device:
 
             h_green = create_green_ctx_handle(
                 c_resources,
-                <unsigned int>n_resources,
-                <cydriver.CUdevice>self._device_id,
-                <unsigned int>cydriver.CUgreenCtxCreate_flags.CU_GREEN_CTX_DEFAULT_STREAM,
+                <unsigned int>(n_resources),
+                <cydriver.CUdevice>(self._device_id),
+                <unsigned int>(cydriver.CUgreenCtxCreate_flags.CU_GREEN_CTX_DEFAULT_STREAM),
             )
             if h_green.get() == NULL:
                 HANDLE_RETURN(get_last_error())
diff --git a/cuda_core/cuda/core/_device_resources.pyx b/cuda_core/cuda/core/_device_resources.pyx
@@ -27,45 +27,85 @@ __all__ = [
 ]
 
 
+# Module-level cached version checks (trinary: 0=unchecked, 1=supported, -1=unsupported)
+cdef int _green_ctx_checked = 0
+cdef int _workqueue_checked = 0
+cdef str _green_ctx_err_msg = ""
+cdef str _workqueue_err_msg = ""
+
+
 cdef inline int _check_green_ctx_support() except?-1:
+    global _green_ctx_checked, _green_ctx_err_msg
+    if _green_ctx_checked == 1:
+        return 0
+    if _green_ctx_checked == -1:
+        raise ValueError(_green_ctx_err_msg)
     cdef tuple drv = cy_driver_version()
     cdef tuple bind = cy_binding_version()
     if drv < (12, 4, 0):
-        raise NotImplementedError(
-            "Green context support requires CUDA driver 12.4 or newer. "
-            f"Using driver version {'.'.join(map(str, drv))}"
+        _green_ctx_err_msg = (
+            "Green context support requires CUDA driver 12.4 or newer "
+            f"(current driver: {'.'.join(map(str, drv))})"
         )
+        _green_ctx_checked = -1
+        raise ValueError(_green_ctx_err_msg)
     if bind < (12, 4, 0):
-        raise NotImplementedError(
-            "Green context support requires cuda.bindings 12.4 or newer. "
-            f"Using cuda.bindings version {'.'.join(map(str, bind))}"
+        _green_ctx_err_msg = (
+            "Green context support requires cuda.bindings 12.4 or newer "
+            f"(current bindings: {'.'.join(map(str, bind))})"
         )
+        _green_ctx_checked = -1
+        raise ValueError(_green_ctx_err_msg)
+    _green_ctx_checked = 1
     return 0
 
 
 cdef inline int _check_workqueue_support() except?-1:
+    global _workqueue_checked, _workqueue_err_msg
+    if _workqueue_checked == 1:
+        return 0
+    if _workqueue_checked == -1:
+        raise ValueError(_workqueue_err_msg)
     cdef tuple drv = cy_driver_version()
     cdef tuple bind = cy_binding_version()
     if drv < (13, 1, 0):
-        raise NotImplementedError(
-            "WorkqueueResource requires CUDA driver 13.1 or newer. "
-            f"Using driver version {'.'.join(map(str, drv))}"
+        _workqueue_err_msg = (
+            "WorkqueueResource requires CUDA driver 13.1 or newer "
+            f"(current driver: {'.'.join(map(str, drv))})"
         )
+        _workqueue_checked = -1
+        raise ValueError(_workqueue_err_msg)
     if bind < (13, 1, 0):
-        raise NotImplementedError(
-            "WorkqueueResource requires cuda.bindings 13.1 or newer. "
-            f"Using cuda.bindings version {'.'.join(map(str, bind))}"
+        _workqueue_err_msg = (
+            "WorkqueueResource requires cuda.bindings 13.1 or newer "
+            f"(current bindings: {'.'.join(map(str, bind))})"
         )
+        _workqueue_checked = -1
+        raise ValueError(_workqueue_err_msg)
+    _workqueue_checked = 1
     return 0
 
 
 @dataclass
 cdef class SMResourceOptions:
-    """Options for :meth:`SMResource.split`.
-
-    ``count`` determines the number of requested groups. Scalar ``count`` or
-    ``None`` creates one group; a sequence creates ``len(count)`` groups. Other
-    sequence fields must match the length of ``count``.
+    """Customizable :obj:`SMResource.split` options.
+
+    Each field accepts a scalar (for a single group) or a ``Sequence``
+    (for multiple groups). ``count`` drives the number of groups; other
+    ``Sequence`` fields must match its length.
+
+    Attributes
+    ----------
+    count : int or Sequence[int], optional
+        Requested SM count per group. ``None`` means discovery mode
+        (auto-detect). (Default to ``None``)
+    coscheduled_sm_count : int or Sequence[int], optional
+        Minimum number of SMs guaranteed to be co-scheduled in each
+        group. (Default to ``None``)
+    preferred_coscheduled_sm_count : int or Sequence[int], optional
+        Preferred co-scheduled SM count; the driver tries to satisfy
+        this but may fall back to ``coscheduled_sm_count``.
+        (Default to ``None``)
     """
 
     count: int | SequenceABC | None = None
@@ -75,7 +115,14 @@ cdef class SMResourceOptions:
 
 @dataclass
 cdef class WorkqueueResourceOptions:
-    """Options for :meth:`WorkqueueResource.configure`."""
+    """Customizable :obj:`WorkqueueResource.configure` options.
+
+    Attributes
+    ----------
+    sharing_scope : str, optional
+        Workqueue sharing scope. Accepted values: ``"device_ctx"``
+        or ``"green_ctx_balanced"``. (Default to ``None``)
+    """
 
     sharing_scope: str | None = None
 
@@ -134,18 +181,13 @@ cdef object _broadcast_field(object value, int n_groups):
     return [value] * n_groups
 
 
-cdef inline unsigned int _as_uint(object value, str field_name) except? 0:
-    if not isinstance(value, int):
-        raise TypeError(f"{field_name} must be an int or None, got {type(value)}")
-    if value < 0:
-        raise ValueError(f"{field_name} must be non-negative")
-    return <unsigned int>value
-
-
-cdef inline unsigned int _count_to_sm_count(object value) except? 0:
+cdef inline unsigned int _to_sm_count(object value) except? 0:
+    """Convert a count value to unsigned int. None maps to 0 (discovery)."""
     if value is None:
         return 0
-    return _as_uint(value, "count")
+    if value < 0:
+        raise ValueError(f"count must be non-negative, got {value}")
+    return <unsigned int>(value)
 
 
 cdef inline bint _can_use_structured_sm_split():
@@ -196,7 +238,7 @@ cdef object _resolve_split_by_count_request(SMResourceOptions options):
                 "use CUDA 13.1 or newer for per-group counts"
             )
 
-    min_count = _count_to_sm_count(first)
+    min_count = _to_sm_count(first)
     return n_groups, min_count
 
 
@@ -213,13 +255,11 @@ IF CUDA_CORE_BUILD_MAJOR >= 13:
 
         for i in range(n_groups):
             memset(&params[i], 0, sizeof(cydriver.CU_DEV_SM_RESOURCE_GROUP_PARAMS))
-            params[i].smCount = _count_to_sm_count(counts[i])
+            params[i].smCount = _to_sm_count(counts[i])
             if coscheduled[i] is not None:
-                params[i].coscheduledSmCount = _as_uint(coscheduled[i], "coscheduled_sm_count")
+                params[i].coscheduledSmCount = <unsigned int>(coscheduled[i])
             if preferred[i] is not None:
-                params[i].preferredCoscheduledSmCount = _as_uint(
-                    preferred[i], "preferred_coscheduled_sm_count"
-                )
+                params[i].preferredCoscheduledSmCount = <unsigned int>(preferred[i])
             params[i].flags = 0
         return 0
 
@@ -253,7 +293,7 @@ IF CUDA_CORE_BUILD_MAJOR >= 13:
             with nogil:
                 HANDLE_RETURN(cydriver.cuDevSmResourceSplit(
                     result,
-                    <unsigned int>n_groups,
+                    <unsigned int>(n_groups),
                     &sm._resource,
                     &remaining,
                     0,
@@ -285,8 +325,8 @@ ELSE:
 
 cdef object _split_with_count_api(SMResource sm, SMResourceOptions options, bint dry_run):
     cdef object request = _resolve_split_by_count_request(options)
-    cdef unsigned int nb_groups = <unsigned int>request[0]
-    cdef unsigned int min_count = <unsigned int>request[1]
+    cdef unsigned int nb_groups = <unsigned int>(request[0])
+    cdef unsigned int min_count = <unsigned int>(request[1])
     cdef unsigned int actual_groups = nb_groups
     cdef cydriver.CUdevResource* result = NULL
     cdef cydriver.CUdevResource remaining
@@ -328,7 +368,7 @@ cdef inline unsigned int _sm_resource_granularity(int device_id) except? 0:
         HANDLE_RETURN(cydriver.cuDeviceGetAttribute(
             &major,
             cydriver.CUdevice_attribute.CU_DEVICE_ATTRIBUTE_COMPUTE_CAPABILITY_MAJOR,
-            <cydriver.CUdevice>device_id,
+            <cydriver.CUdevice>(device_id),
         ))
     if major >= 9:
         return 8
@@ -342,7 +382,11 @@ cdef inline unsigned int _fallback_if_zero(unsigned int value, unsigned int fall
 
 
 cdef class SMResource:
-    """SM resource queried from a device. Not user-constructible."""
+    """Represent an SM (streaming multiprocessor) resource partition.
+
+    Instances are returned by :obj:`DeviceResources.sm` or
+    :meth:`SMResource.split` and cannot be instantiated directly.
+    """
 
     def __init__(self, *args, **kwargs):
         raise RuntimeError(
@@ -391,7 +435,7 @@ cdef class SMResource:
     @property
     def handle(self) -> int:
         """Return the address of the underlying ``CUdevResource`` struct."""
-        return <intptr_t>&self._resource
+        return <intptr_t>(&self._resource)
 
     @property
     def sm_count(self) -> int:
@@ -414,7 +458,22 @@ cdef class SMResource:
         return self._flags
 
     def split(self, options not None, *, bint dry_run=False):
-        """Split this SM resource into groups plus a remainder."""
+        """Split this SM resource into groups and a remainder.
+
+        Parameters
+        ----------
+        options : :obj:`SMResourceOptions`
+            Split configuration (count, co-scheduling constraints).
+        dry_run : bool, optional
+            If ``True``, return filled-in metadata without creating
+            usable resource objects. (Default to ``False``)
+
+        Returns
+        -------
+        tuple[list[:obj:`SMResource`], :obj:`SMResource`]
+            ``(groups, remainder)`` where each group holds a disjoint
+            SM partition and *remainder* holds any unassigned SMs.
+        """
         cdef SMResourceOptions opts = check_or_create_options(
             SMResourceOptions, options, "SM resource options"
         )
@@ -427,7 +486,13 @@ cdef class SMResource:
 
 
 cdef class WorkqueueResource:
-    """Workqueue resource. Not user-constructible."""
+    """Represent a workqueue resource for a device or green context.
+
+    Merges ``CU_DEV_RESOURCE_TYPE_WORKQUEUE_CONFIG`` and
+    ``CU_DEV_RESOURCE_TYPE_WORKQUEUE`` under one user-facing type.
+    Instances are returned by :obj:`DeviceResources.workqueue` and
+    cannot be instantiated directly.
+    """
 
     def __init__(self, *args, **kwargs):
         raise RuntimeError(
@@ -448,10 +513,16 @@ cdef class WorkqueueResource:
     @property
     def handle(self) -> int:
         """Return the address of the underlying config ``CUdevResource`` struct."""
-        return <intptr_t>&self._wq_config_resource
+        return <intptr_t>(&self._wq_config_resource)
 
     def configure(self, options not None):
-        """Configure the workqueue resource in place."""
+        """Configure the workqueue resource in place.
+
+        Parameters
+        ----------
+        options : :obj:`WorkqueueResourceOptions`
+            Configuration options (sharing scope, etc.).
+        """
         cdef WorkqueueResourceOptions opts = check_or_create_options(
             WorkqueueResourceOptions, options, "Workqueue resource options"
         )
@@ -481,11 +552,14 @@ cdef class WorkqueueResource:
 
 
 cdef class DeviceResources:
-    """Namespace for hardware resource query. Not user-constructible.
+    """Namespace for hardware resource queries.
+
+    When obtained via :obj:`Device.resources`, queries return full device
+    resources. When obtained via :obj:`Context.resources` or
+    :obj:`Stream.resources`, queries return the resources provisioned for
+    that context.
 
-    When obtained via ``dev.resources``, queries return full device resources.
-    When obtained via ``ctx.resources``, queries return the resources
-    provisioned for that context.
+    This class cannot be instantiated directly.
     """
 
     def __init__(self, *args, **kwargs):
@@ -525,14 +599,14 @@ cdef class DeviceResources:
                 ))
         else:
             HANDLE_RETURN(cydriver.cuDeviceGetDevResource(
-                <cydriver.CUdevice>self._device_id, res,
+                <cydriver.CUdevice>(self._device_id), res,
                 cydriver.CUdevResourceType.CU_DEV_RESOURCE_TYPE_SM,
             ))
         return 0
 
     @property
     def sm(self) -> SMResource:
-        """Query SM resources."""
+        """Return the :obj:`SMResource` for this device or context."""
         _check_green_ctx_support()
         cdef cydriver.CUdevResource res
         with nogil:
@@ -541,7 +615,7 @@ cdef class DeviceResources:
 
     @property
     def workqueue(self) -> WorkqueueResource:
-        """Query workqueue resources."""
+        """Return the :obj:`WorkqueueResource` for this device or context."""
         _check_green_ctx_support()
         _check_workqueue_support()
         cdef cydriver.CUdevResource _wq_config
@@ -581,12 +655,12 @@ cdef class DeviceResources:
                 # Device-level query
                 with nogil:
                     HANDLE_RETURN(cydriver.cuDeviceGetDevResource(
-                        <cydriver.CUdevice>self._device_id,
+                        <cydriver.CUdevice>(self._device_id),
                         &_wq_config,
                         cydriver.CUdevResourceType.CU_DEV_RESOURCE_TYPE_WORKQUEUE_CONFIG,
                     ))
                     HANDLE_RETURN(cydriver.cuDeviceGetDevResource(
-                        <cydriver.CUdevice>self._device_id,
+                        <cydriver.CUdevice>(self._device_id),
                         &_wq,
                         cydriver.CUdevResourceType.CU_DEV_RESOURCE_TYPE_WORKQUEUE,
                     ))
diff --git a/cuda_core/docs/source/api.rst b/cuda_core/docs/source/api.rst
@@ -26,12 +26,18 @@ Devices and execution
 
    Stream
    Event
+   Context
+   SMResource
+   WorkqueueResource
 
    :template: dataclass.rst
 
    StreamOptions
    EventOptions
    LaunchConfig
+   ContextOptions
+   SMResourceOptions
+   WorkqueueResourceOptions
 
 .. data:: LEGACY_DEFAULT_STREAM
 
diff --git a/cuda_core/docs/source/api_private.rst b/cuda_core/docs/source/api_private.rst
@@ -31,6 +31,7 @@ CUDA runtime
    :template: autosummary/cyclass.rst
 
    _device.DeviceProperties
+   _device_resources.DeviceResources
    _memory._ipc.IPCAllocationHandle
    _memory._ipc.IPCBufferDescriptor
 
diff --git a/cuda_core/tests/test_green_context.py b/cuda_core/tests/test_green_context.py
