|
19 | 19 | # cython: linetrace=True |
20 | 20 |
|
21 | 21 | import logging |
22 | | -from contextlib import ExitStack, contextmanager |
23 | 22 | from contextvars import ContextVar |
24 | | - |
25 | | -from .enum_types import backend_type, device_type |
26 | | - |
27 | | -from ._backend cimport ( # noqa: E211 |
28 | | - DPCTLQueueMgr_GetCurrentQueue, |
29 | | - DPCTLQueueMgr_GetQueueStackSize, |
30 | | - DPCTLQueueMgr_GlobalQueueIsCurrent, |
31 | | - DPCTLQueueMgr_PopQueue, |
32 | | - DPCTLQueueMgr_PushQueue, |
33 | | - DPCTLQueueMgr_SetGlobalQueue, |
34 | | - DPCTLSyclQueueRef, |
35 | | - _backend_type, |
36 | | - _device_type, |
37 | | -) |
38 | 23 | from ._sycl_context cimport SyclContext |
39 | 24 | from ._sycl_device cimport SyclDevice |
40 | 25 |
|
41 | 26 | __all__ = [ |
42 | 27 | "_global_device_queue_cache", |
43 | | - "device_context", |
44 | | - "get_current_backend", |
45 | | - "get_current_device_type", |
46 | | - "get_current_queue", |
47 | 28 | "get_device_cached_queue", |
48 | | - "get_num_activated_queues", |
49 | | - "is_in_device_context", |
50 | | - "set_global_queue", |
51 | 29 | ] |
52 | 30 |
|
53 | 31 | _logger = logging.getLogger(__name__) |
54 | 32 |
|
55 | 33 |
|
56 | | -cdef class _SyclQueueManager: |
57 | | - """ Provides a SYCL queue manager interface for Python. |
58 | | - """ |
59 | | - |
60 | | - def _set_as_current_queue(self, arg): |
61 | | - cdef SyclQueue q |
62 | | - cdef DPCTLSyclQueueRef queue_ref = NULL |
63 | | - |
64 | | - if isinstance(arg, SyclQueue): |
65 | | - q_obj = arg |
66 | | - else: |
67 | | - q_obj = SyclQueue(arg) |
68 | | - |
69 | | - q = <SyclQueue> q_obj |
70 | | - queue_ref = q.get_queue_ref() |
71 | | - DPCTLQueueMgr_PushQueue(queue_ref) |
72 | | - |
73 | | - return q_obj |
74 | | - |
75 | | - def _remove_current_queue(self): |
76 | | - DPCTLQueueMgr_PopQueue() |
77 | | - |
78 | | - cpdef get_current_backend(self): |
79 | | - """ |
80 | | - Returns the backend for the current queue as a `backend_type` enum. |
81 | | -
|
82 | | - Returns: |
83 | | - backend_type: The SYCL backend for the currently selected queue. |
84 | | - """ |
85 | | - return self.get_current_queue().backend |
86 | | - |
87 | | - cpdef get_current_device_type(self): |
88 | | - """ |
89 | | - Returns current device type as a `device_type` enum. |
90 | | -
|
91 | | - Returns: |
92 | | - device_type: The SYCL device type for the currently selected queue. |
93 | | - Possible values can be gpu, cpu, accelerator, or host. |
94 | | - """ |
95 | | - return self.get_current_queue().sycl_device.device_type |
96 | | - |
97 | | - cpdef SyclQueue get_current_queue(self): |
98 | | - """ |
99 | | - Returns the currently activated SYCL queue as a new SyclQueue object. |
100 | | -
|
101 | | - Returns: |
102 | | - SyclQueue: If there is a currently active SYCL queue that queue |
103 | | - is returned wrapped in a SyclQueue object. The SyclQueue object |
104 | | - owns a copy of the currently active SYCL queue as an opaque |
105 | | - `DPCTLSyclQueueRef` pointer. The pointer is freed when the SyclQueue |
106 | | - is garbage collected. |
107 | | -
|
108 | | - Raises: |
109 | | - SyclQueueCreationError: If no currently active SYCL queue found. |
110 | | - """ |
111 | | - return SyclQueue._create(DPCTLQueueMgr_GetCurrentQueue()) |
112 | | - |
113 | | - def get_num_activated_queues(self): |
114 | | - """ |
115 | | - Returns the number of currently activated queues for this thread. |
116 | | -
|
117 | | - Whenever a program's control enters a :func:`dpctl.device_context()` |
118 | | - scope, either a new SYCL queue is created or a previously created |
119 | | - queue is retrieved from a cache and yielded. The queue yielded by the |
120 | | - context manager is termed to be "activated". If a program creates |
121 | | - multiple nested :func:`dpctl.device_context()` scopes then multiple |
122 | | - queues can be activated at the same time, although only the latest |
123 | | - activated queue is usable directly via calling |
124 | | - :func:`dpctl.get_current_queue()`. This function returns the number of |
125 | | - currently activated queues. |
126 | | -
|
127 | | - Returns: |
128 | | - int: The number of currently activated queues. |
129 | | -
|
130 | | - """ |
131 | | - return DPCTLQueueMgr_GetQueueStackSize() |
132 | | - |
133 | | - def is_in_device_context(self): |
134 | | - """ |
135 | | - Checks if the control is inside a :func:`dpctl.device_context()` scope. |
136 | | -
|
137 | | - Returns: |
138 | | - bool: True if the control is within a |
139 | | - :func:`dpctl.device_context()` scope, otherwise False. |
140 | | - """ |
141 | | - cdef int inCtx = DPCTLQueueMgr_GlobalQueueIsCurrent() |
142 | | - return not bool(inCtx) |
143 | | - |
144 | | - def set_global_queue(self, arg): |
145 | | - """ |
146 | | - Sets the global queue to the SYCL queue specified explicitly, |
147 | | - or created from given arguments. |
148 | | -
|
149 | | - Args: |
150 | | - arg: An instance of :class:`dpctl.SyclQueue` or a filter selector |
151 | | - string to be used to construct a :class:`dpctl.SyclQueue`. The |
152 | | - queue is stored in the dpctl queue manager as the default queue. |
153 | | - Raises: |
154 | | - SyclQueueCreationError: If a SYCL queue could not be created. |
155 | | - """ |
156 | | - cdef SyclQueue q |
157 | | - cdef DPCTLSyclQueueRef queue_ref = NULL |
158 | | - |
159 | | - if type(arg) is SyclQueue: |
160 | | - q = <SyclQueue> arg |
161 | | - else: |
162 | | - q_obj = SyclQueue(arg) |
163 | | - q = <SyclQueue> q_obj |
164 | | - |
165 | | - queue_ref = q.get_queue_ref() |
166 | | - DPCTLQueueMgr_SetGlobalQueue(queue_ref) |
167 | | - |
168 | | - |
169 | | -# This private instance of the _SyclQueueManager should not be directly |
170 | | -# accessed outside the module. |
171 | | -_mgr = _SyclQueueManager() |
172 | | - |
173 | | -# Global bound functions |
174 | | -get_num_activated_queues = _mgr.get_num_activated_queues |
175 | | -set_global_queue = _mgr.set_global_queue |
176 | | -is_in_device_context = _mgr.is_in_device_context |
177 | | - |
178 | | - |
179 | | -cpdef SyclQueue get_current_queue(): |
180 | | - """ |
181 | | - Returns the currently activate SYCL queue as a new SyclQueue object. |
182 | | -
|
183 | | - Returns: |
184 | | - SyclQueue: If there is a currently active SYCL queue that queue |
185 | | - is returned wrapped in a SyclQueue object. The SyclQueue object |
186 | | - owns a copy of the currently active SYCL queue as an opaque |
187 | | - `DPCTLSyclQueueRef` pointer. The pointer is freed when the SyclQueue |
188 | | - is garbage collected. |
189 | | -
|
190 | | - Raises: |
191 | | - SyclQueueCreationError: If no currently active SYCL queue found. |
192 | | - """ |
193 | | - return _mgr.get_current_queue() |
194 | | - |
195 | | - |
196 | | -cpdef get_current_device_type(): |
197 | | - """ |
198 | | - Returns current device type as a `device_type` enum. |
199 | | -
|
200 | | - Returns: |
201 | | - device_type: The SYCL device type for the currently selected queue. |
202 | | - Possible values can be gpu, cpu, accelerator, or host. |
203 | | - """ |
204 | | - return _mgr.get_current_device_type() |
205 | | - |
206 | | - |
207 | | -cpdef get_current_backend(): |
208 | | - """ |
209 | | - Returns the backend for the current queue as a `backend_type` enum. |
210 | | -
|
211 | | - Returns: |
212 | | - backend_type: The SYCL backend for the currently selected queue. |
213 | | - """ |
214 | | - return _mgr.get_current_backend() |
215 | | - |
216 | | - |
217 | | -nested_context_factories = [] |
218 | | - |
219 | | - |
220 | | -def _get_nested_contexts(ctxt): |
221 | | - _help_numba_dppy() |
222 | | - return (factory(ctxt) for factory in nested_context_factories) |
223 | | - |
224 | | - |
225 | | -def _help_numba_dppy(): |
226 | | - """Import numba-dppy for registering nested contexts""" |
227 | | - try: |
228 | | - import numba_dppy |
229 | | - except Exception: |
230 | | - pass |
231 | | - |
232 | | - |
233 | | -@contextmanager |
234 | | -def device_context(arg): |
235 | | - """ |
236 | | - Yields a SYCL queue corresponding to the input queue object, device object, |
237 | | - or device filter selector string. |
238 | | -
|
239 | | - This context manager "activates", *i.e.*, sets as the currently usable |
240 | | - queue, the SYCL queue defined by the argument `arg`. |
241 | | - The activated queue is yielded by the context manager and can also be |
242 | | - accessed by any subsequent call to :func:`dpctl.get_current_queue()` inside |
243 | | - the context manager's scope. The yielded queue is removed as the currently |
244 | | - usable queue on exiting the context manager. |
245 | | -
|
246 | | - You can register context factory in the list of factories. |
247 | | - This context manager uses context factories to create and activate nested contexts. |
248 | | -
|
249 | | - Args: |
250 | | - arg : A :class:`dpctl.SyclQueue` object, or a :class:`dpctl.SyclDevice` |
251 | | - object, or a filter selector string. |
252 | | -
|
253 | | - Yields: |
254 | | - :class:`dpctl.SyclQueue`: A SYCL queue corresponding to the specified |
255 | | - input device, queue, or filter string. |
256 | | -
|
257 | | - Raises: |
258 | | - SyclQueueCreationError: If the SYCL queue creation failed. |
259 | | -
|
260 | | - :Example: |
261 | | - The following example sets current queue targeting specific device |
262 | | - indicated with filter selector string in the scope of `with` block: |
263 | | -
|
264 | | - .. code-block:: python |
265 | | -
|
266 | | - import dpctl |
267 | | - with dpctl.device_context("level0:gpu:0"): |
268 | | - do_something_on_gpu0() |
269 | | -
|
270 | | - The following example registers nested context factory: |
271 | | -
|
272 | | - .. code-block:: python |
273 | | -
|
274 | | - import dctl |
275 | | -
|
276 | | - def factory(sycl_queue): |
277 | | - ... |
278 | | - return context |
279 | | -
|
280 | | - dpctl.nested_context_factories.append(factory) |
281 | | -
|
282 | | - """ |
283 | | - ctxt = None |
284 | | - try: |
285 | | - ctxt = _mgr._set_as_current_queue(arg) |
286 | | - with ExitStack() as stack: |
287 | | - for nested_context in _get_nested_contexts(ctxt): |
288 | | - stack.enter_context(nested_context) |
289 | | - yield ctxt |
290 | | - finally: |
291 | | - # Code to release resource |
292 | | - if ctxt: |
293 | | - _logger.debug( |
294 | | - "Removing the queue from the stack of active queues") |
295 | | - _mgr._remove_current_queue() |
296 | | - else: |
297 | | - _logger.debug("No queue was created so nothing to do") |
298 | | - |
299 | | - |
300 | 34 | cdef class _DeviceDefaultQueueCache: |
301 | 35 | cdef dict __device_queue_map__ |
302 | 36 |
|
|
0 commit comments