Skip to content

Commit 08d52d1

Browse files
leofangclaude
leofang
and
claude
committed
Add explicit green context model: ctx.create_stream and ctx.resources
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). - Delegates to Stream._init, which calls create_stream_handle in C++. - C++ create_stream_handle auto-dispatches: checks get_context_green_ctx and calls cuGreenCtxStreamCreate for green contexts, or cuStreamCreateWithPriority for primary. Single function, no duplication. Context.resources: - Returns a DeviceResources namespace querying this context's resources (cuCtxGetDevResource / cuGreenCtxGetDevResource), not the full device. 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 regression tests with _use_green_ctx helper. Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
1 parent aee88c6 commit 08d52d1

8 files changed

Lines changed: 214 additions & 65 deletions

File tree

‎cuda_core/cuda/core/_context.pyx‎

Lines changed: 43 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -8,6 +8,7 @@ from collections.abc import Sequence
88
from dataclasses import dataclass
99

1010
from cuda.bindings cimport cydriver
11+
from cuda.core._device_resources cimport DeviceResources, SMResource, WorkqueueResource
1112
from cuda.core._device_resources import SMResource, WorkqueueResource
1213
from cuda.core._resource_handles cimport (
1314
ContextHandle,
@@ -19,6 +20,7 @@ from cuda.core._resource_handles cimport (
1920
as_intptr,
2021
as_py,
2122
)
23+
from cuda.core._stream import StreamOptions
2224
from cuda.core._utils.cuda_utils cimport HANDLE_RETURN
2325

2426

@@ -75,6 +77,47 @@ cdef class Context:
7577
return False
7678
return get_context_green_ctx(self._h_context).get() != NULL
7779

80+
@property
81+
def resources(self) -> DeviceResources:
82+
"""Query the hardware resources provisioned for this context.
83+
84+
For green contexts, returns the resources this context was created
85+
with (SM partition, workqueue config). For primary contexts, returns
86+
the full device resources.
87+
88+
Raises :class:`RuntimeError` if the context has been closed.
89+
"""
90+
if not self._h_context:
91+
raise RuntimeError("Cannot query resources on a closed context")
92+
return DeviceResources._init_from_ctx(self._h_context, self._device_id)
93+
94+
def create_stream(self, options: StreamOptions | None = None):
95+
"""Create a new stream bound to this green context.
96+
97+
This method is only available on green contexts. For primary
98+
contexts, use :meth:`Device.create_stream` instead.
99+
100+
Parameters
101+
----------
102+
options : :obj:`~_stream.StreamOptions`, optional
103+
Customizable dataclass for stream creation options.
104+
105+
Returns
106+
-------
107+
:obj:`~_stream.Stream`
108+
Newly created stream object.
109+
"""
110+
if not self._h_context:
111+
raise RuntimeError("Cannot create a stream on a closed context")
112+
if not self.is_green:
113+
raise RuntimeError(
114+
"Context.create_stream() is only supported on green contexts. "
115+
"Use Device.create_stream() for primary contexts."
116+
)
117+
118+
from cuda.core._stream import Stream
119+
return Stream._init(options=options, device_id=self._device_id, ctx=self)
120+
78121
cpdef close(self):
79122
"""Release this context wrapper's underlying CUDA handles."""
80123
cdef cydriver.CUcontext current_ctx

‎cuda_core/cuda/core/_cpp/resource_handles.cpp‎

Lines changed: 13 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -34,6 +34,8 @@ decltype(&cuGreenCtxDestroy) p_cuGreenCtxDestroy = nullptr;
3434
decltype(&cuCtxFromGreenCtx) p_cuCtxFromGreenCtx = nullptr;
3535
decltype(&cuDevResourceGenerateDesc) p_cuDevResourceGenerateDesc = nullptr;
3636

37+
decltype(&cuGreenCtxStreamCreate) p_cuGreenCtxStreamCreate = nullptr;
38+
3739
decltype(&cuStreamCreateWithPriority) p_cuStreamCreateWithPriority = nullptr;
3840
decltype(&cuStreamDestroy) p_cuStreamDestroy = nullptr;
3941

@@ -409,8 +411,17 @@ static HandleRegistry<CUstream, StreamHandle> stream_registry;
409411
StreamHandle create_stream_handle(const ContextHandle& h_ctx, unsigned int flags, int priority) {
410412
GILReleaseGuard gil;
411413
CUstream stream;
412-
if (CUDA_SUCCESS != (err = p_cuStreamCreateWithPriority(&stream, flags, priority))) {
413-
return {};
414+
415+
// Dispatch: green context uses cuGreenCtxStreamCreate, primary uses cuStreamCreateWithPriority
416+
GreenCtxHandle h_green = get_context_green_ctx(h_ctx);
417+
if (h_green && p_cuGreenCtxStreamCreate) {
418+
if (CUDA_SUCCESS != (err = p_cuGreenCtxStreamCreate(&stream, as_cu(h_green), flags, priority))) {
419+
return {};
420+
}
421+
} else {
422+
if (CUDA_SUCCESS != (err = p_cuStreamCreateWithPriority(&stream, flags, priority))) {
423+
return {};
424+
}
414425
}
415426

416427
auto box = std::shared_ptr<const StreamBox>(

‎cuda_core/cuda/core/_cpp/resource_handles.hpp‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -64,6 +64,8 @@ extern decltype(&cuGreenCtxDestroy) p_cuGreenCtxDestroy;
6464
extern decltype(&cuCtxFromGreenCtx) p_cuCtxFromGreenCtx;
6565
extern decltype(&cuDevResourceGenerateDesc) p_cuDevResourceGenerateDesc;
6666

67+
extern decltype(&cuGreenCtxStreamCreate) p_cuGreenCtxStreamCreate;
68+
6769
extern decltype(&cuStreamCreateWithPriority) p_cuStreamCreateWithPriority;
6870
extern decltype(&cuStreamDestroy) p_cuStreamDestroy;
6971

‎cuda_core/cuda/core/_device_resources.pxd‎

Lines changed: 7 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -3,6 +3,7 @@
33
# SPDX-License-Identifier: Apache-2.0
44

55
from cuda.bindings cimport cydriver
6+
from cuda.core._resource_handles cimport ContextHandle, GreenCtxHandle
67

78

89
cdef class SMResource:
@@ -38,7 +39,13 @@ cdef class WorkqueueResource:
3839
cdef class DeviceResources:
3940
cdef:
4041
int _device_id
42+
ContextHandle _h_context # NULL for device-level queries
4143
object __weakref__
4244

4345
@staticmethod
4446
cdef DeviceResources _init(int device_id)
47+
48+
@staticmethod
49+
cdef DeviceResources _init_from_ctx(ContextHandle h_context, int device_id)
50+
51+
cdef inline int _query_sm(self, cydriver.CUdevResource* res) except?-1 nogil

‎cuda_core/cuda/core/_device_resources.pyx‎

Lines changed: 67 additions & 20 deletions
Original file line numberDiff line numberDiff line change
@@ -12,6 +12,7 @@ from libc.stdlib cimport free, malloc
1212
from libc.string cimport memset
1313

1414
from cuda.bindings cimport cydriver
15+
from cuda.core._resource_handles cimport ContextHandle, GreenCtxHandle, as_cu, get_context_green_ctx
1516
from cuda.core._utils.cuda_utils cimport check_or_create_options, HANDLE_RETURN
1617
from cuda.core._utils.cuda_utils import is_sequence
1718
from cuda.core._utils.version cimport cy_binding_version, cy_driver_version
@@ -480,53 +481,99 @@ cdef class WorkqueueResource:
480481

481482

482483
cdef class DeviceResources:
483-
"""Namespace for hardware resource query. Not user-constructible."""
484+
"""Namespace for hardware resource query. Not user-constructible.
485+
486+
When obtained via ``dev.resources``, queries return full device resources.
487+
When obtained via ``ctx.resources``, queries return the resources
488+
provisioned for that context.
489+
"""
484490

485491
def __init__(self, *args, **kwargs):
486492
raise RuntimeError(
487493
"DeviceResources cannot be instantiated directly. "
488-
"Use dev.resources."
494+
"Use dev.resources or ctx.resources."
489495
)
490496

491497
@staticmethod
492498
cdef DeviceResources _init(int device_id):
493499
cdef DeviceResources self = DeviceResources.__new__(DeviceResources)
494500
self._device_id = device_id
501+
# _h_context is default empty — queries use cuDeviceGetDevResource
495502
return self
496503

504+
@staticmethod
505+
cdef DeviceResources _init_from_ctx(ContextHandle h_context, int device_id):
506+
cdef DeviceResources self = DeviceResources.__new__(DeviceResources)
507+
self._device_id = device_id
508+
self._h_context = h_context
509+
return self
510+
511+
cdef inline int _query_sm(self, cydriver.CUdevResource* res) except?-1 nogil:
512+
"""Query SM resource from either device or context."""
513+
cdef GreenCtxHandle h_green
514+
if self._h_context:
515+
h_green = get_context_green_ctx(self._h_context)
516+
if h_green:
517+
HANDLE_RETURN(cydriver.cuGreenCtxGetDevResource(
518+
as_cu(h_green), res,
519+
cydriver.CUdevResourceType.CU_DEV_RESOURCE_TYPE_SM,
520+
))
521+
else:
522+
HANDLE_RETURN(cydriver.cuCtxGetDevResource(
523+
as_cu(self._h_context), res,
524+
cydriver.CUdevResourceType.CU_DEV_RESOURCE_TYPE_SM,
525+
))
526+
else:
527+
HANDLE_RETURN(cydriver.cuDeviceGetDevResource(
528+
<cydriver.CUdevice>self._device_id, res,
529+
cydriver.CUdevResourceType.CU_DEV_RESOURCE_TYPE_SM,
530+
))
531+
return 0
532+
497533
@property
498534
def sm(self) -> SMResource:
499-
"""Query SM resources from this device."""
535+
"""Query SM resources."""
500536
_check_green_ctx_support()
501537
cdef cydriver.CUdevResource res
502538
with nogil:
503-
HANDLE_RETURN(cydriver.cuDeviceGetDevResource(
504-
<cydriver.CUdevice>self._device_id,
505-
&res,
506-
cydriver.CUdevResourceType.CU_DEV_RESOURCE_TYPE_SM,
507-
))
539+
self._query_sm(&res)
508540
return SMResource._from_dev_resource(res, self._device_id)
509541

510542
@property
511543
def workqueue(self) -> WorkqueueResource:
512-
"""Query workqueue resources from this device."""
544+
"""Query workqueue resources."""
513545
_check_green_ctx_support()
514546
_check_workqueue_support()
515547
cdef cydriver.CUdevResource _wq_config
516548
cdef cydriver.CUdevResource _wq
517549

518550
IF CUDA_CORE_BUILD_MAJOR >= 13:
519-
with nogil:
520-
HANDLE_RETURN(cydriver.cuDeviceGetDevResource(
521-
<cydriver.CUdevice>self._device_id,
522-
&_wq_config,
523-
cydriver.CUdevResourceType.CU_DEV_RESOURCE_TYPE_WORKQUEUE_CONFIG,
524-
))
525-
HANDLE_RETURN(cydriver.cuDeviceGetDevResource(
526-
<cydriver.CUdevice>self._device_id,
527-
&_wq,
528-
cydriver.CUdevResourceType.CU_DEV_RESOURCE_TYPE_WORKQUEUE,
529-
))
551+
if self._h_context:
552+
# Context-level query
553+
with nogil:
554+
HANDLE_RETURN(cydriver.cuCtxGetDevResource(
555+
as_cu(self._h_context),
556+
&_wq_config,
557+
cydriver.CUdevResourceType.CU_DEV_RESOURCE_TYPE_WORKQUEUE_CONFIG,
558+
))
559+
HANDLE_RETURN(cydriver.cuCtxGetDevResource(
560+
as_cu(self._h_context),
561+
&_wq,
562+
cydriver.CUdevResourceType.CU_DEV_RESOURCE_TYPE_WORKQUEUE,
563+
))
564+
else:
565+
# Device-level query
566+
with nogil:
567+
HANDLE_RETURN(cydriver.cuDeviceGetDevResource(
568+
<cydriver.CUdevice>self._device_id,
569+
&_wq_config,
570+
cydriver.CUdevResourceType.CU_DEV_RESOURCE_TYPE_WORKQUEUE_CONFIG,
571+
))
572+
HANDLE_RETURN(cydriver.cuDeviceGetDevResource(
573+
<cydriver.CUdevice>self._device_id,
574+
&_wq,
575+
cydriver.CUdevResourceType.CU_DEV_RESOURCE_TYPE_WORKQUEUE,
576+
))
530577
return WorkqueueResource._from_dev_resources(_wq_config, _wq)
531578
ELSE:
532579
raise NotImplementedError(

‎cuda_core/cuda/core/_resource_handles.pyx‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -239,6 +239,7 @@ cdef extern from "_cpp/resource_handles.hpp" namespace "cuda_core":
239239
void* p_cuGreenCtxDestroy "reinterpret_cast<void*&>(cuda_core::p_cuGreenCtxDestroy)"
240240
void* p_cuCtxFromGreenCtx "reinterpret_cast<void*&>(cuda_core::p_cuCtxFromGreenCtx)"
241241
void* p_cuDevResourceGenerateDesc "reinterpret_cast<void*&>(cuda_core::p_cuDevResourceGenerateDesc)"
242+
void* p_cuGreenCtxStreamCreate "reinterpret_cast<void*&>(cuda_core::p_cuGreenCtxStreamCreate)"
242243

243244
# Stream
244245
void* p_cuStreamCreateWithPriority "reinterpret_cast<void*&>(cuda_core::p_cuStreamCreateWithPriority)"
@@ -320,6 +321,7 @@ p_cuGreenCtxCreate = _get_optional_driver_fn("cuGreenCtxCreate")
320321
p_cuGreenCtxDestroy = _get_optional_driver_fn("cuGreenCtxDestroy")
321322
p_cuCtxFromGreenCtx = _get_optional_driver_fn("cuCtxFromGreenCtx")
322323
p_cuDevResourceGenerateDesc = _get_optional_driver_fn("cuDevResourceGenerateDesc")
324+
p_cuGreenCtxStreamCreate = _get_optional_driver_fn("cuGreenCtxStreamCreate")
323325

324326
# Stream
325327
p_cuStreamCreateWithPriority = _get_driver_fn("cuStreamCreateWithPriority")

‎cuda_core/cuda/core/_stream.pyx‎

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -153,7 +153,8 @@ cdef class Stream:
153153
else:
154154
prio = high
155155

156-
# C++ creates the stream and returns owning handle with context dependency
156+
# C++ creates the stream and returns owning handle with context dependency.
157+
# For green contexts, the C++ layer auto-dispatches to cuGreenCtxStreamCreate.
157158
h_stream = create_stream_handle(h_context, flags, prio)
158159
if not h_stream:
159160
raise RuntimeError("Failed to create CUDA stream")

0 commit comments

Comments
 (0)