Add green context support

cuda_core/cuda/core/__init__.py

@@ -29,7 +29,15 @@ def _import_versioned_module():
 
 
 from cuda.core import checkpoint, 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)