make SequentialOrderManager thread-local and cached queues, devices global

Author: ndgrigorian

dpctl/_sycl_device_factory.pyx

@@ -19,7 +19,8 @@
 # cython: linetrace=True
 # cython: freethreading_compatible = True
 
-""" This module implements several device creation helper functions:
+"""
+This module implements several device creation helper functions:
 
   - wrapper functions to create a SyclDevice from the standard SYCL
     device selector classes.
@@ -47,7 +48,7 @@ from ._backend cimport (  # noqa: E211
     _device_type,
 )
 
-from contextvars import ContextVar
+import threading
 
 from ._sycl_device import SyclDeviceCreationError
 from .enum_types import backend_type
@@ -287,7 +288,8 @@ cpdef int get_num_devices(
 
 
 cpdef cpp_bool has_cpu_devices():
-    """ A helper function to check if there are any SYCL CPU devices available.
+    """
+    A helper function to check if there are any SYCL CPU devices available.
 
     Returns:
         bool:
@@ -299,7 +301,8 @@ cpdef cpp_bool has_cpu_devices():
 
 
 cpdef cpp_bool has_gpu_devices():
-    """ A helper function to check if there are any SYCL GPU devices available.
+    """
+    A helper function to check if there are any SYCL GPU devices available.
 
     Returns:
         bool:
@@ -311,7 +314,8 @@ cpdef cpp_bool has_gpu_devices():
 
 
 cpdef cpp_bool has_accelerator_devices():
-    """ A helper function to check if there are any SYCL Accelerator devices
+    """
+    A helper function to check if there are any SYCL Accelerator devices
     available.
 
     Returns:
@@ -326,7 +330,8 @@ cpdef cpp_bool has_accelerator_devices():
 
 
 cpdef SyclDevice select_accelerator_device():
-    """A wrapper for ``sycl::device{sycl::accelerator_selector_v}`` constructor.
+    """
+    A wrapper for ``sycl::device{sycl::accelerator_selector_v}`` constructor.
 
     Returns:
         dpctl.SyclDevice:
@@ -348,7 +353,8 @@ cpdef SyclDevice select_accelerator_device():
 
 
 cpdef SyclDevice select_cpu_device():
-    """A wrapper for ``sycl::device{sycl::cpu_selector_v}`` constructor.
+    """
+    A wrapper for ``sycl::device{sycl::cpu_selector_v}`` constructor.
 
     Returns:
         dpctl.SyclDevice:
@@ -370,7 +376,8 @@ cpdef SyclDevice select_cpu_device():
 
 
 cpdef SyclDevice select_default_device():
-    """A wrapper for ``sycl::device{sycl::default_selector_v}`` constructor.
+    """
+    A wrapper for ``sycl::device{sycl::default_selector_v}`` constructor.
 
     Returns:
         dpctl.SyclDevice:
@@ -392,7 +399,8 @@ cpdef SyclDevice select_default_device():
 
 
 cpdef SyclDevice select_gpu_device():
-    """A wrapper for ``sycl::device{sycl::gpu_selector_v}`` constructor.
+    """
+    A wrapper for ``sycl::device{sycl::gpu_selector_v}`` constructor.
 
     Returns:
         dpctl.SyclDevice:
@@ -415,21 +423,23 @@ cpdef SyclDevice select_gpu_device():
 
 cdef class _DefaultDeviceCache:
     cdef dict __device_map__
+    cdef object _cache_lock
 
     def __cinit__(self):
         self.__device_map__ = dict()
-
-    cdef get_or_create(self):
-        """Return instance of SyclDevice and indicator if cache
-        has been modified"""
-        key = 0
-        if key in self.__device_map__:
-            return self.__device_map__[key], False
-        dev = select_default_device()
-        self.__device_map__[key] = dev
-        return dev, True
-
-    cdef _update_map(self, dev_map):
+        self._cache_lock = threading.Lock()
+
+    def get_or_create(self):
+        """Return cached default SyclDevice, creating it if needed."""
+        with self._cache_lock:
+            key = 0
+            if key in self.__device_map__:
+                return self.__device_map__[key]
+            dev = select_default_device()
+            self.__device_map__[key] = dev
+            return dev
+
+    def _update_map(self, dev_map):
         self.__device_map__.update(dev_map)
 
     def __copy__(self):
@@ -439,37 +449,17 @@ cdef class _DefaultDeviceCache:
         return _copy
 
 
-# no default, as would share a single mutable instance across threads and
-# concurrent access to the cache would not be thread-safe. Using ContextVar
-# without a default ensures each context gets its own instance.
-_global_default_device_cache = ContextVar(
-    "global_default_device_cache",
-)
-
-
-cdef _DefaultDeviceCache _get_default_device_cache():
-    """
-    Factory function to get or create a default device cache for the current
-    context
-    """
-    try:
-        return _global_default_device_cache.get()
-    except LookupError:
-        cache = _DefaultDeviceCache()
-        _global_default_device_cache.set(cache)
-        return cache
+# all threads share the same cached default
+_global_default_device_cache = _DefaultDeviceCache()
 
 
 cpdef SyclDevice _cached_default_device():
-    """Returns a cached device selected by default selector.
+    """
+    Returns a cached device selected by default selector.
 
     Returns:
         dpctl.SyclDevice:
             A cached default-selected SYCL device.
 
     """
-    cdef _DefaultDeviceCache _cache = _get_default_device_cache()
-    d_, changed_ = _cache.get_or_create()
-    if changed_:
-        _global_default_device_cache.set(_cache)
-    return d_
+    return _global_default_device_cache.get_or_create()

dpctl/_sycl_queue_manager.pyx

@@ -20,7 +20,7 @@
 # cython: freethreading_compatible = True
 
 import logging
-from contextvars import ContextVar
+import threading
 from ._sycl_context cimport SyclContext
 from ._sycl_device cimport SyclDevice
 
@@ -34,13 +34,14 @@ _logger = logging.getLogger(__name__)
 
 cdef class _DeviceDefaultQueueCache:
     cdef dict __device_queue_map__
+    cdef object _cache_lock
 
     def __cinit__(self):
         self.__device_queue_map__ = dict()
+        self._cache_lock = threading.Lock()
 
     def get_or_create(self, key):
-        """Return instance of SyclQueue and indicator if cache
-        has been modified"""
+        """Return cached SyclQueue for given key, creating it if needed."""
         if (
             isinstance(key, tuple)
             and len(key) == 2
@@ -57,14 +58,15 @@ cdef class _DeviceDefaultQueueCache:
             ctx_dev = q.sycl_context, q.sycl_device
         else:
             raise TypeError
-        if ctx_dev in self.__device_queue_map__:
-            return self.__device_queue_map__[ctx_dev], False
-        if q is None:
-            q = SyclQueue(*ctx_dev)
-        self.__device_queue_map__[ctx_dev] = q
-        return q, True
-
-    cdef _update_map(self, dev_queue_map):
+        with self._cache_lock:
+            if ctx_dev in self.__device_queue_map__:
+                return self.__device_queue_map__[ctx_dev]
+            if q is None:
+                q = SyclQueue(*ctx_dev)
+            self.__device_queue_map__[ctx_dev] = q
+            return q
+
+    def _update_map(self, dev_queue_map):
         self.__device_queue_map__.update(dev_queue_map)
 
     def __copy__(self):
@@ -75,29 +77,13 @@ cdef class _DeviceDefaultQueueCache:
         return _copy
 
 
-# no default, as would share a single mutable instance across threads and
-# concurrent access to the cache would not be thread-safe. Using ContextVar
-# without a default ensures each context gets its own instance.
-_global_device_queue_cache = ContextVar(
-    "global_device_queue_cache",
-)
-
-
-cdef _DeviceDefaultQueueCache _get_device_queue_cache():
-    """
-    Factory function to get or create a default device queue cache for the
-    current context
-    """
-    try:
-        return _global_device_queue_cache.get()
-    except LookupError:
-        cache = _DeviceDefaultQueueCache()
-        _global_device_queue_cache.set(cache)
-        return cache
+# all threads share the same cached default
+_global_device_queue_cache = _DeviceDefaultQueueCache()
 
 
 cpdef object get_device_cached_queue(object key):
-    """Returns a cached queue associated with given device.
+    """
+    Returns a cached queue associated with given device.
 
     Args:
         key : Either a 2-tuple consisting of a :class:`dpctl.SyclContext` and
@@ -112,8 +98,4 @@ cpdef object get_device_cached_queue(object key):
         TypeError: If the input key is not one of the accepted types.
 
     """
-    _cache = _get_device_queue_cache()
-    q_, changed_ = _cache.get_or_create(key)
-    if changed_:
-        _global_device_queue_cache.set(_cache)
-    return q_
+    return _global_device_queue_cache.get_or_create(key)

dpctl/tests/test_sycl_queue_manager.py

@@ -24,19 +24,20 @@
 def test__DeviceDefaultQueueCache():
     import copy
 
-    from dpctl._sycl_queue_manager import _global_device_queue_cache as cache
-    from dpctl._sycl_queue_manager import get_device_cached_queue
+    from dpctl._sycl_queue_manager import (
+        _global_device_queue_cache,
+        get_device_cached_queue,
+    )
 
     try:
         d = dpctl.SyclDevice()
     except dpctl.SyclDeviceCreationError:
         pytest.skip("Could not create default device")
 
     q1 = get_device_cached_queue(d)
-    cache_copy = copy.copy(cache.get())
-    q2, changed = cache_copy.get_or_create(d)
+    cache_copy = copy.copy(_global_device_queue_cache)
+    q2 = cache_copy.get_or_create(d)
 
-    assert not changed
     assert q1 == q2
     q3 = get_device_cached_queue(d.filter_string)
     assert q3 == q1

dpctl/utils/_order_manager.py

@@ -1,6 +1,6 @@
+import threading
 import weakref
 from collections import defaultdict
-from contextvars import ContextVar
 
 from .._sycl_event import SyclEvent
 from .._sycl_queue import SyclQueue
@@ -65,26 +65,28 @@ def __copy__(self):
 
 
 class SyclQueueToOrderManagerMap:
-    """Utility class used to ensure sequential ordering of offloaded tasks
-    when passed to order manager."""
+    """
+    Utility class used to ensure sequential ordering of offloaded tasks
+    when passed to order manager.
+
+    Maintains a thread-local dictionary mapping SyclQueue instances to
+    _SequentialOrderManager instances.
+    """
 
     def __init__(self):
-        self._map = ContextVar(
-            "global_order_manager_map",
-            # no default to avoid sharing a single defaultdict
-            # across threads
-        )
+        # each thread gets its own dictionary of order managers
+        self._tls = threading.local()
 
     def _get_map(self):
         """
-        Factory method to get or create a default device queue cache for the
-        current context
+        Factory method to get or create an order manager map for the
+        current thread.
         """
         try:
-            return self._map.get()
-        except LookupError:
+            return self._tls.order_manager_map
+        except AttributeError:
             m = defaultdict(_SequentialOrderManager)
-            self._map.set(m)
+            self._tls.order_manager_map = m
             return m
 
     def __getitem__(self, q: SyclQueue) -> _SequentialOrderManager: