Skip to content

Commit b88ebc6

Browse files
committed
Merge remote-tracking branch 'origin/main' into explicit-stream-2001
Co-authored-by: Cursor <cursoragent@cursor.com> # Conflicts: # cuda_core/cuda/core/_memory/_buffer.pyx # cuda_core/cuda/core/_memory/_graph_memory_resource.pyx # cuda_core/cuda/core/_memory/_legacy.py # cuda_core/cuda/core/_memory/_memory_pool.pyx # cuda_core/cuda/core/_module.pyx # cuda_core/docs/source/release/1.0.0-notes.rst
2 parents e1d3b8c + cec0efb commit b88ebc6

44 files changed

Lines changed: 2932 additions & 254 deletions

Some content is hidden

Large Commits have some content hidden by default. Use the searchbox below for content that may be hidden.

cuda_core/cuda/core/__init__.py

Lines changed: 9 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -28,8 +28,16 @@ def _import_versioned_module():
2828
del _import_versioned_module
2929

3030

31-
from cuda.core import system, utils
31+
from cuda.core import checkpoint, system, utils
32+
from cuda.core._context import Context, ContextOptions
3233
from cuda.core._device import Device
34+
from cuda.core._device_resources import (
35+
DeviceResources,
36+
SMResource,
37+
SMResourceOptions,
38+
WorkqueueResource,
39+
WorkqueueResourceOptions,
40+
)
3341
from cuda.core._event import Event, EventOptions
3442
from cuda.core._graphics import GraphicsResource
3543
from cuda.core._launch_config import LaunchConfig

cuda_core/cuda/core/_context.pxd

Lines changed: 7 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1,8 +1,8 @@
1-
# SPDX-FileCopyrightText: Copyright (c) 2024-2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
1+
# SPDX-FileCopyrightText: Copyright (c) 2024-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
22
#
33
# SPDX-License-Identifier: Apache-2.0
44

5-
from cuda.core._resource_handles cimport ContextHandle
5+
from cuda.core._resource_handles cimport ContextHandle, GreenCtxHandle
66

77
cdef class Context:
88
"""Cython declaration for Context class.
@@ -18,3 +18,8 @@ cdef class Context:
1818

1919
@staticmethod
2020
cdef Context _from_handle(type cls, ContextHandle h_context, int device_id)
21+
22+
@staticmethod
23+
cdef Context _from_green_ctx(type cls, GreenCtxHandle h_green_ctx, int device_id)
24+
25+
cpdef close(self)

cuda_core/cuda/core/_context.pyx

Lines changed: 94 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -2,18 +2,34 @@
22
#
33
# SPDX-License-Identifier: Apache-2.0
44

5+
from __future__ import annotations
6+
7+
from collections.abc import Sequence
58
from dataclasses import dataclass
69

10+
from cuda.bindings cimport cydriver
11+
from cuda.core._device_resources cimport DeviceResources, SMResource, WorkqueueResource
12+
from cuda.core._device_resources import SMResource, WorkqueueResource
713
from cuda.core._resource_handles cimport (
814
ContextHandle,
15+
GreenCtxHandle,
16+
as_cu,
17+
create_context_handle_from_green_ctx,
18+
get_context_green_ctx,
19+
get_last_error,
920
as_intptr,
1021
as_py,
1122
)
23+
from cuda.core._stream import Stream, StreamOptions
24+
from cuda.core._utils.cuda_utils cimport HANDLE_RETURN
1225

1326

1427
__all__ = ['Context', 'ContextOptions']
1528

1629

30+
DeviceResourcesType = Sequence[SMResource | WorkqueueResource]
31+
32+
1733
cdef class Context:
1834
"""CUDA context wrapper.
1935
@@ -32,17 +48,88 @@ cdef class Context:
3248
ctx._device_id = device_id
3349
return ctx
3450

51+
@staticmethod
52+
cdef Context _from_green_ctx(type cls, GreenCtxHandle h_green_ctx, int device_id):
53+
"""Create Context from an owning green context handle."""
54+
cdef ContextHandle h_context = create_context_handle_from_green_ctx(h_green_ctx)
55+
if not h_context:
56+
HANDLE_RETURN(get_last_error())
57+
raise RuntimeError("Failed to create CUDA context view from green context")
58+
return Context._from_handle(cls, h_context, device_id)
59+
3560
@property
3661
def handle(self):
3762
"""Return the underlying CUcontext handle."""
38-
if self._h_context.get() == NULL:
63+
if not self._h_context:
64+
return None
65+
if as_cu(self._h_context) == NULL:
3966
return None
4067
return as_py(self._h_context)
4168

4269
@property
4370
def _handle(self):
4471
return self.handle
4572

73+
@property
74+
def is_green(self) -> bool:
75+
"""True if this context was created from device resources."""
76+
if not self._h_context:
77+
return False
78+
return get_context_green_ctx(self._h_context).get() != NULL
79+
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+
return Stream._init(options=options, device_id=self._device_id, ctx=self)
119+
120+
cpdef close(self):
121+
"""Release this context wrapper's underlying CUDA handles."""
122+
cdef cydriver.CUcontext current_ctx
123+
if self._h_context and as_cu(self._h_context) != NULL:
124+
with nogil:
125+
HANDLE_RETURN(cydriver.cuCtxGetCurrent(&current_ctx))
126+
if current_ctx == as_cu(self._h_context):
127+
raise RuntimeError(
128+
"Cannot close a CUDA context while it is current. "
129+
"Restore a previous context before closing this context."
130+
)
131+
self._h_context.reset()
132+
46133
def __eq__(self, other):
47134
if not isinstance(other, Context):
48135
return NotImplemented
@@ -57,9 +144,12 @@ cdef class Context:
57144

58145

59146
@dataclass
60-
class ContextOptions:
147+
cdef class ContextOptions:
61148
"""Options for context creation.
62149
63-
Currently unused, reserved for future use.
150+
Attributes
151+
----------
152+
resources : :obj:`~cuda.core.typing.DeviceResourcesType`
153+
Device resources used to create a green context.
64154
"""
65-
pass # TODO
155+
resources: DeviceResourcesType

cuda_core/cuda/core/_cpp/REGISTRY_DESIGN.md

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -29,7 +29,8 @@ carries timing/IPC flags, `KernelBox` carries the library dependency).
2929
Without this level, a round-tripped handle would produce a new Box
3030
with default metadata, losing information that was set at creation.
3131

32-
Instances: `event_registry`, `kernel_registry`, `graph_node_registry`.
32+
Instances: `context_registry`, `stream_registry`, `event_registry`,
33+
`kernel_registry`, `graph_node_registry`.
3334

3435
## Level 2: Resource Handle -> Python Object (Cython)
3536

0 commit comments

Comments
 (0)