Add green context support

cuda_core/cuda/core/__init__.py

@@ -29,7 +29,15 @@ def _import_versioned_module():
 
 
 from cuda.core import system, utils
+from cuda.core._context import Context, ContextOptions
 from cuda.core._device import Device
+from cuda.core._device_resources import (
+    DeviceResources,
+    SMResource,
+    SMResourceOptions,
+    WorkqueueResource,
+    WorkqueueResourceOptions,
+)
 from cuda.core._event import Event, EventOptions
 from cuda.core._graphics import GraphicsResource
 from cuda.core._launch_config import LaunchConfig

cuda_core/cuda/core/_context.pxd

@@ -2,7 +2,7 @@
 #
 # SPDX-License-Identifier: Apache-2.0
 
-from cuda.core._resource_handles cimport ContextHandle
+from cuda.core._resource_handles cimport ContextHandle, GreenCtxHandle
 
 cdef class Context:
     """Cython declaration for Context class.
@@ -18,3 +18,8 @@ cdef class Context:
 
     @staticmethod
     cdef Context _from_handle(type cls, ContextHandle h_context, int device_id)
+
+    @staticmethod
+    cdef Context _from_green_ctx(type cls, GreenCtxHandle h_green_ctx, int device_id)
+
+    cpdef close(self)

cuda_core/cuda/core/_context.pyx

@@ -2,18 +2,34 @@
 #
 # SPDX-License-Identifier: Apache-2.0
 
+from __future__ import annotations
+
+from collections.abc import Sequence
 from dataclasses import dataclass
 
+from cuda.bindings cimport cydriver
+from cuda.core._device_resources cimport DeviceResources, SMResource, WorkqueueResource
+from cuda.core._device_resources import SMResource, WorkqueueResource
 from cuda.core._resource_handles cimport (
     ContextHandle,
+    GreenCtxHandle,
+    as_cu,
+    create_context_handle_from_green_ctx,
+    get_context_green_ctx,
+    get_last_error,
     as_intptr,
     as_py,
 )
+from cuda.core._stream import Stream, StreamOptions
+from cuda.core._utils.cuda_utils cimport HANDLE_RETURN
 
 
 __all__ = ['Context', 'ContextOptions']
 
 
+DeviceResourcesT = Sequence[SMResource | WorkqueueResource]
+
+
 cdef class Context:
     """CUDA context wrapper.
 
@@ -32,17 +48,88 @@ cdef class Context:
         ctx._device_id = device_id
         return ctx
 
+    @staticmethod
+    cdef Context _from_green_ctx(type cls, GreenCtxHandle h_green_ctx, int device_id):
+        """Create Context from an owning green context handle."""
+        cdef ContextHandle h_context = create_context_handle_from_green_ctx(h_green_ctx)
+        if not h_context:
+            HANDLE_RETURN(get_last_error())
+            raise RuntimeError("Failed to create CUDA context view from green context")
+        return Context._from_handle(cls, h_context, device_id)
+
     @property
     def handle(self):
         """Return the underlying CUcontext handle."""
-        if self._h_context.get() == NULL:
+        if not self._h_context:
+            return None
+        if as_cu(self._h_context) == NULL:
             return None
         return as_py(self._h_context)
 
     @property
     def _handle(self):
         return self.handle
 
+    @property
+    def is_green(self) -> bool:
+        """True if this context was created from device resources."""
+        if not self._h_context:
+            return False
+        return get_context_green_ctx(self._h_context).get() != NULL
+
+    @property
+    def resources(self) -> DeviceResources:
+        """Query the hardware resources provisioned for this context.
+
+        For green contexts, returns the resources this context was created
+        with (SM partition, workqueue config). For primary contexts, returns
+        the full device resources.
+
+        Raises :class:`RuntimeError` if the context has been closed.
+        """
+        if not self._h_context:
+            raise RuntimeError("Cannot query resources on a closed context")
+        return DeviceResources._init_from_ctx(self._h_context, self._device_id)
+
+    def create_stream(self, options: StreamOptions | None = None):
+        """Create a new stream bound to this green context.
+
+        This method is only available on green contexts. For primary
+        contexts, use :meth:`Device.create_stream` instead.
+
+        Parameters
+        ----------
+        options : :obj:`~_stream.StreamOptions`, optional
+            Customizable dataclass for stream creation options.
+
+        Returns
+        -------
+        :obj:`~_stream.Stream`
+            Newly created stream object.
+        """
+        if not self._h_context:
+            raise RuntimeError("Cannot create a stream on a closed context")
+        if not self.is_green:
+            raise RuntimeError(
+                "Context.create_stream() is only supported on green contexts. "
+                "Use Device.create_stream() for primary contexts."
+            )
+
+        return Stream._init(options=options, device_id=self._device_id, ctx=self)
+
+    cpdef close(self):
+        """Release this context wrapper's underlying CUDA handles."""
+        cdef cydriver.CUcontext current_ctx
+        if self._h_context and as_cu(self._h_context) != NULL:
+            with nogil:
+                HANDLE_RETURN(cydriver.cuCtxGetCurrent(&current_ctx))
+            if current_ctx == as_cu(self._h_context):
+                raise RuntimeError(
+                    "Cannot close a CUDA context while it is current. "
+                    "Restore a previous context before closing this context."
+                )
+        self._h_context.reset()
+
     def __eq__(self, other):
         if not isinstance(other, Context):
             return NotImplemented
@@ -57,9 +144,12 @@ cdef class Context:
 
 
 @dataclass
-class ContextOptions:
+cdef class ContextOptions:
     """Options for context creation.
 
-    Currently unused, reserved for future use.
+    Attributes
+    ----------
+    resources : :obj:`~cuda.core.typing.DeviceResourcesT`
+        Device resources used to create a green context.
     """
-    pass  # TODO
+    resources: DeviceResourcesT

cuda_core/cuda/core/_cpp/REGISTRY_DESIGN.md

@@ -29,7 +29,8 @@ carries timing/IPC flags, `KernelBox` carries the library dependency).
 Without this level, a round-tripped handle would produce a new Box
 with default metadata, losing information that was set at creation.
 
-Instances: `event_registry`, `kernel_registry`, `graph_node_registry`.
+Instances: `context_registry`, `stream_registry`, `event_registry`,
+`kernel_registry`, `graph_node_registry`.
 
 ## Level 2: Resource Handle -> Python Object (Cython)