Add explicit green context model: ctx.create_stream and ctx.resources

leofang · claude · leofang · commit f7ac69f83eea · 2026-04-27T02:27:39.000Z
Switch from the v1.0 push model (dev.set_current + dev.create_stream)
to the v1.1 explicit model (ctx.create_stream + ctx.resources) as the
primary way to use green contexts.

Context.create_stream(options):
- Only supported on green contexts (raises on primary contexts).
- Calls cuGreenCtxStreamCreate via the C++ handle layer — no push/pop
  of the context stack needed.
- New C++ create_stream_handle_for_green_ctx() with function pointer
  for cuGreenCtxStreamCreate.

Context.resources:
- Returns a DeviceResources namespace querying this context's resources
  (cuCtxGetDevResource / cuGreenCtxGetDevResource), not the full device.
- DeviceResources._init_from_ctx() uses the ContextHandle to dispatch.

dev.set_current(green_ctx) still works but is not the recommended path.

Tests rewritten to use the explicit model throughout. Push-model
set_current kept as a single regression test.

Co-Authored-By: Claude Opus 4.6 (1M context) &lt;noreply@anthropic.com&gt;
diff --git a/cuda_core/cuda/core/_context.pyx b/cuda_core/cuda/core/_context.pyx
@@ -8,18 +8,23 @@ 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,
+    StreamHandle,
     as_cu,
     create_context_handle_from_green_ctx,
+    create_stream_handle_for_green_ctx,
     get_context_green_ctx,
     get_last_error,
     as_intptr,
     as_py,
 )
-from cuda.core._utils.cuda_utils cimport HANDLE_RETURN
+from cuda.core._stream cimport Stream as cyStream
+from cuda.core._stream import StreamOptions
+from cuda.core._utils.cuda_utils cimport check_or_create_options, HANDLE_RETURN
 
 
 __all__ = ['Context', 'ContextOptions']
@@ -75,6 +80,75 @@ cdef class 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."
+            )
+
+        opts = check_or_create_options(StreamOptions, options, "Stream options")
+        cdef bint nonblocking = opts.nonblocking
+        cdef int priority = 0
+        cdef int high, low
+        cdef unsigned int flags
+
+        flags = (cydriver.CUstream_flags.CU_STREAM_NON_BLOCKING if nonblocking
+                 else cydriver.CUstream_flags.CU_STREAM_DEFAULT)
+
+        with nogil:
+            HANDLE_RETURN(cydriver.cuCtxGetStreamPriorityRange(&high, &low))
+        if opts.priority is not None:
+            priority = opts.priority
+            if not (low <= priority <= high):
+                raise ValueError(f"priority={priority} is out of range [{low}, {high}]")
+        else:
+            priority = high
+
+        cdef GreenCtxHandle h_green = get_context_green_ctx(self._h_context)
+        cdef StreamHandle h_stream = create_stream_handle_for_green_ctx(
+            h_green, self._h_context, flags, priority)
+        if not h_stream:
+            HANDLE_RETURN(get_last_error())
+            raise RuntimeError("Failed to create CUDA stream for green context")
+
+        cdef cyStream s = cyStream._from_handle(cyStream, h_stream)
+        s._nonblocking = int(nonblocking)
+        s._priority = priority
+        s._device_id = self._device_id
+        return s
+
     cpdef close(self):
         """Release this context wrapper's underlying CUDA handles."""
         cdef cydriver.CUcontext current_ctx
diff --git a/cuda_core/cuda/core/_cpp/resource_handles.cpp b/cuda_core/cuda/core/_cpp/resource_handles.cpp
@@ -34,6 +34,8 @@ decltype(&cuGreenCtxDestroy) p_cuGreenCtxDestroy = nullptr;
 decltype(&cuCtxFromGreenCtx) p_cuCtxFromGreenCtx = nullptr;
 decltype(&cuDevResourceGenerateDesc) p_cuDevResourceGenerateDesc = nullptr;
 
+decltype(&cuGreenCtxStreamCreate) p_cuGreenCtxStreamCreate = nullptr;
+
 decltype(&cuStreamCreateWithPriority) p_cuStreamCreateWithPriority = nullptr;
 decltype(&cuStreamDestroy) p_cuStreamDestroy = nullptr;
 
@@ -427,6 +429,34 @@ StreamHandle create_stream_handle(const ContextHandle& h_ctx, unsigned int flags
     return h;
 }
 
+StreamHandle create_stream_handle_for_green_ctx(
+    const GreenCtxHandle& h_green, const ContextHandle& h_ctx,
+    unsigned int flags, int priority) {
+    if (!p_cuGreenCtxStreamCreate) {
+        err = CUDA_ERROR_NOT_SUPPORTED;
+        return {};
+    }
+
+    GILReleaseGuard gil;
+    CUstream stream;
+    if (CUDA_SUCCESS != (err = p_cuGreenCtxStreamCreate(&stream, as_cu(h_green), flags, priority))) {
+        return {};
+    }
+
+    auto box = std::shared_ptr<const StreamBox>(
+        new StreamBox{stream, h_ctx},
+        [](const StreamBox* b) {
+            stream_registry.unregister_handle(b->resource);
+            GILReleaseGuard gil;
+            p_cuStreamDestroy(b->resource);
+            delete b;
+        }
+    );
+    StreamHandle h(box, &box->resource);
+    stream_registry.register_handle(stream, h);
+    return h;
+}
+
 StreamHandle create_stream_handle_ref(CUstream stream) {
     if (auto h = stream_registry.lookup(stream)) {
         return h;
diff --git a/cuda_core/cuda/core/_cpp/resource_handles.hpp b/cuda_core/cuda/core/_cpp/resource_handles.hpp
@@ -64,6 +64,8 @@ extern decltype(&cuGreenCtxDestroy) p_cuGreenCtxDestroy;
 extern decltype(&cuCtxFromGreenCtx) p_cuCtxFromGreenCtx;
 extern decltype(&cuDevResourceGenerateDesc) p_cuDevResourceGenerateDesc;
 
+extern decltype(&cuGreenCtxStreamCreate) p_cuGreenCtxStreamCreate;
+
 extern decltype(&cuStreamCreateWithPriority) p_cuStreamCreateWithPriority;
 extern decltype(&cuStreamDestroy) p_cuStreamDestroy;
 
@@ -213,6 +215,12 @@ StreamHandle create_stream_handle_ref(CUstream stream);
 // The owner is responsible for keeping the stream's context alive.
 StreamHandle create_stream_handle_with_owner(CUstream stream, PyObject* owner);
 
+// Create an owning stream from a green context (calls cuGreenCtxStreamCreate).
+// The returned stream handle depends on the green ctx's ContextHandle.
+StreamHandle create_stream_handle_for_green_ctx(
+    const GreenCtxHandle& h_green, const ContextHandle& h_ctx,
+    unsigned int flags, int priority);
+
 // Return the context dependency associated with a stream handle, if any.
 ContextHandle get_stream_context(const StreamHandle& h) noexcept;
 
diff --git a/cuda_core/cuda/core/_device_resources.pxd b/cuda_core/cuda/core/_device_resources.pxd
@@ -3,6 +3,7 @@
 # SPDX-License-Identifier: Apache-2.0
 
 from cuda.bindings cimport cydriver
+from cuda.core._resource_handles cimport ContextHandle, GreenCtxHandle
 
 
 cdef class SMResource:
@@ -38,7 +39,13 @@ cdef class WorkqueueResource:
 cdef class DeviceResources:
     cdef:
         int _device_id
+        ContextHandle _h_context  # NULL for device-level queries
         object __weakref__
 
     @staticmethod
     cdef DeviceResources _init(int device_id)
+
+    @staticmethod
+    cdef DeviceResources _init_from_ctx(ContextHandle h_context, int device_id)
+
+    cdef inline int _query_sm(self, cydriver.CUdevResource* res) except?-1 nogil
diff --git a/cuda_core/cuda/core/_device_resources.pyx b/cuda_core/cuda/core/_device_resources.pyx
@@ -12,6 +12,7 @@ from libc.stdlib cimport free, malloc
 from libc.string cimport memset
 
 from cuda.bindings cimport cydriver
+from cuda.core._resource_handles cimport ContextHandle, GreenCtxHandle, as_cu, get_context_green_ctx
 from cuda.core._utils.cuda_utils cimport check_or_create_options, HANDLE_RETURN
 from cuda.core._utils.cuda_utils import is_sequence
 from cuda.core._utils.version cimport cy_binding_version, cy_driver_version
@@ -480,53 +481,99 @@ cdef class WorkqueueResource:
 
 
 cdef class DeviceResources:
-    """Namespace for hardware resource query. Not user-constructible."""
+    """Namespace for hardware resource query. Not user-constructible.
+
+    When obtained via ``dev.resources``, queries return full device resources.
+    When obtained via ``ctx.resources``, queries return the resources
+    provisioned for that context.
+    """
 
     def __init__(self, *args, **kwargs):
         raise RuntimeError(
             "DeviceResources cannot be instantiated directly. "
-            "Use dev.resources."
+            "Use dev.resources or ctx.resources."
         )
 
     @staticmethod
     cdef DeviceResources _init(int device_id):
         cdef DeviceResources self = DeviceResources.__new__(DeviceResources)
         self._device_id = device_id
+        # _h_context is default empty — queries use cuDeviceGetDevResource
         return self
 
+    @staticmethod
+    cdef DeviceResources _init_from_ctx(ContextHandle h_context, int device_id):
+        cdef DeviceResources self = DeviceResources.__new__(DeviceResources)
+        self._device_id = device_id
+        self._h_context = h_context
+        return self
+
+    cdef inline int _query_sm(self, cydriver.CUdevResource* res) except?-1 nogil:
+        """Query SM resource from either device or context."""
+        cdef GreenCtxHandle h_green
+        if self._h_context:
+            h_green = get_context_green_ctx(self._h_context)
+            if h_green:
+                HANDLE_RETURN(cydriver.cuGreenCtxGetDevResource(
+                    as_cu(h_green), res,
+                    cydriver.CUdevResourceType.CU_DEV_RESOURCE_TYPE_SM,
+                ))
+            else:
+                HANDLE_RETURN(cydriver.cuCtxGetDevResource(
+                    as_cu(self._h_context), res,
+                    cydriver.CUdevResourceType.CU_DEV_RESOURCE_TYPE_SM,
+                ))
+        else:
+            HANDLE_RETURN(cydriver.cuDeviceGetDevResource(
+                <cydriver.CUdevice>self._device_id, res,
+                cydriver.CUdevResourceType.CU_DEV_RESOURCE_TYPE_SM,
+            ))
+        return 0
+
     @property
     def sm(self) -> SMResource:
-        """Query SM resources from this device."""
+        """Query SM resources."""
         _check_green_ctx_support()
         cdef cydriver.CUdevResource res
         with nogil:
-            HANDLE_RETURN(cydriver.cuDeviceGetDevResource(
-                <cydriver.CUdevice>self._device_id,
-                &res,
-                cydriver.CUdevResourceType.CU_DEV_RESOURCE_TYPE_SM,
-            ))
+            self._query_sm(&res)
         return SMResource._from_dev_resource(res, self._device_id)
 
     @property
     def workqueue(self) -> WorkqueueResource:
-        """Query workqueue resources from this device."""
+        """Query workqueue resources."""
         _check_green_ctx_support()
         _check_workqueue_support()
         cdef cydriver.CUdevResource _wq_config
         cdef cydriver.CUdevResource _wq
 
         IF CUDA_CORE_BUILD_MAJOR >= 13:
-            with nogil:
-                HANDLE_RETURN(cydriver.cuDeviceGetDevResource(
-                    <cydriver.CUdevice>self._device_id,
-                    &_wq_config,
-                    cydriver.CUdevResourceType.CU_DEV_RESOURCE_TYPE_WORKQUEUE_CONFIG,
-                ))
-                HANDLE_RETURN(cydriver.cuDeviceGetDevResource(
-                    <cydriver.CUdevice>self._device_id,
-                    &_wq,
-                    cydriver.CUdevResourceType.CU_DEV_RESOURCE_TYPE_WORKQUEUE,
-                ))
+            if self._h_context:
+                # Context-level query
+                with nogil:
+                    HANDLE_RETURN(cydriver.cuCtxGetDevResource(
+                        as_cu(self._h_context),
+                        &_wq_config,
+                        cydriver.CUdevResourceType.CU_DEV_RESOURCE_TYPE_WORKQUEUE_CONFIG,
+                    ))
+                    HANDLE_RETURN(cydriver.cuCtxGetDevResource(
+                        as_cu(self._h_context),
+                        &_wq,
+                        cydriver.CUdevResourceType.CU_DEV_RESOURCE_TYPE_WORKQUEUE,
+                    ))
+            else:
+                # Device-level query
+                with nogil:
+                    HANDLE_RETURN(cydriver.cuDeviceGetDevResource(
+                        <cydriver.CUdevice>self._device_id,
+                        &_wq_config,
+                        cydriver.CUdevResourceType.CU_DEV_RESOURCE_TYPE_WORKQUEUE_CONFIG,
+                    ))
+                    HANDLE_RETURN(cydriver.cuDeviceGetDevResource(
+                        <cydriver.CUdevice>self._device_id,
+                        &_wq,
+                        cydriver.CUdevResourceType.CU_DEV_RESOURCE_TYPE_WORKQUEUE,
+                    ))
             return WorkqueueResource._from_dev_resources(_wq_config, _wq)
         ELSE:
             raise NotImplementedError(
diff --git a/cuda_core/cuda/core/_resource_handles.pxd b/cuda_core/cuda/core/_resource_handles.pxd
@@ -123,6 +123,9 @@ cdef ContextHandle get_current_context() except+ nogil
 # Stream handles
 cdef StreamHandle create_stream_handle(
     const ContextHandle& h_ctx, unsigned int flags, int priority) except+ nogil
+cdef StreamHandle create_stream_handle_for_green_ctx(
+    const GreenCtxHandle& h_green, const ContextHandle& h_ctx,
+    unsigned int flags, int priority) except+ nogil
 cdef StreamHandle create_stream_handle_ref(cydriver.CUstream stream) except+ nogil
 cdef StreamHandle create_stream_handle_with_owner(cydriver.CUstream stream, object owner) except+ nogil
 cdef ContextHandle get_stream_context(const StreamHandle& h) noexcept nogil
diff --git a/cuda_core/cuda/core/_resource_handles.pyx b/cuda_core/cuda/core/_resource_handles.pyx
@@ -72,6 +72,9 @@ cdef extern from "_cpp/resource_handles.hpp" namespace "cuda_core":
     # Stream handles
     StreamHandle create_stream_handle "cuda_core::create_stream_handle" (
         const ContextHandle& h_ctx, unsigned int flags, int priority) except+ nogil
+    StreamHandle create_stream_handle_for_green_ctx "cuda_core::create_stream_handle_for_green_ctx" (
+        const GreenCtxHandle& h_green, const ContextHandle& h_ctx,
+        unsigned int flags, int priority) except+ nogil
     StreamHandle create_stream_handle_ref "cuda_core::create_stream_handle_ref" (
         cydriver.CUstream stream) except+ nogil
     StreamHandle create_stream_handle_with_owner "cuda_core::create_stream_handle_with_owner" (
@@ -239,6 +242,7 @@ cdef extern from "_cpp/resource_handles.hpp" namespace "cuda_core":
     void* p_cuGreenCtxDestroy "reinterpret_cast<void*&>(cuda_core::p_cuGreenCtxDestroy)"
     void* p_cuCtxFromGreenCtx "reinterpret_cast<void*&>(cuda_core::p_cuCtxFromGreenCtx)"
     void* p_cuDevResourceGenerateDesc "reinterpret_cast<void*&>(cuda_core::p_cuDevResourceGenerateDesc)"
+    void* p_cuGreenCtxStreamCreate "reinterpret_cast<void*&>(cuda_core::p_cuGreenCtxStreamCreate)"
 
     # Stream
     void* p_cuStreamCreateWithPriority "reinterpret_cast<void*&>(cuda_core::p_cuStreamCreateWithPriority)"
@@ -320,6 +324,7 @@ p_cuGreenCtxCreate = _get_optional_driver_fn("cuGreenCtxCreate")
 p_cuGreenCtxDestroy = _get_optional_driver_fn("cuGreenCtxDestroy")
 p_cuCtxFromGreenCtx = _get_optional_driver_fn("cuCtxFromGreenCtx")
 p_cuDevResourceGenerateDesc = _get_optional_driver_fn("cuDevResourceGenerateDesc")
+p_cuGreenCtxStreamCreate = _get_optional_driver_fn("cuGreenCtxStreamCreate")
 
 # Stream
 p_cuStreamCreateWithPriority = _get_driver_fn("cuStreamCreateWithPriority")
diff --git a/cuda_core/tests/test_green_context.py b/cuda_core/tests/test_green_context.py
