22#
33# SPDX-License-Identifier: Apache-2.0
44
5+ from __future__ import annotations
6+
7+ from collections.abc import Sequence
58from 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
713from 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+
1733cdef 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
0 commit comments