Make ordermanager free-threading safe

Author: ndgrigorian

dpctl/_sycl_queue.pxd

@@ -43,6 +43,7 @@ cdef public api class _SyclQueue [
     cdef DPCTLSyclQueueRef _queue_ref
     cdef SyclContext _context
     cdef SyclDevice _device
+    cdef object __weakref__
 
 
 cdef public api class SyclQueue (_SyclQueue) [

dpctl/utils/_order_manager.py

@@ -16,7 +16,7 @@ class _SequentialOrderManager:
     def __init__(self):
         self._state = _OrderManager(16)
 
-    def __dealloc__(self):
+    def __del__(self):
         _local = self._state
         SyclEvent.wait_for(_local.get_submitted_events())
         SyclEvent.wait_for(_local.get_host_task_events())
@@ -71,24 +71,31 @@ class SyclQueueToOrderManagerMap:
     def __init__(self):
         self._map = ContextVar(
             "global_order_manager_map",
-            default=defaultdict(_SequentialOrderManager),
+            # no default to avoid sharing a single defaultdict
+            # across threads
         )
 
+    def _get_map(self):
+        """
+        Factory method to get or create a default device queue cache for the
+        current context
+        """
+        try:
+            return self._map.get()
+        except LookupError:
+            m = defaultdict(_SequentialOrderManager)
+            self._map.set(m)
+            return m
+
     def __getitem__(self, q: SyclQueue) -> _SequentialOrderManager:
         """Get order manager for given SyclQueue"""
-        _local = self._map.get()
         if not isinstance(q, SyclQueue):
             raise TypeError(f"Expected `dpctl.SyclQueue`, got {type(q)}")
-        if q in _local:
-            return _local[q]
-        else:
-            v = _local[q]
-            _local[q] = v
-            return v
+        return self._get_map()[q]
 
     def clear(self):
         """Clear content of internal dictionary"""
-        _local = self._map.get()
+        _local = self._get_map()
         for v in _local.values():
             v.wait()
         _local.clear()

dpctl/utils/src/sequential_order_keeper.hpp

@@ -3,6 +3,7 @@
 
 #include <algorithm>
 #include <cstddef>
+#include <mutex>
 #include <vector>
 
 namespace
@@ -21,10 +22,12 @@ inline bool is_event_complete(const sycl::event &e)
 class SequentialOrder
 {
 private:
+    mutable std::mutex mu_events;
     std::vector<sycl::event> host_task_events;
     std::vector<sycl::event> submitted_events;
 
-    void prune_complete()
+    // only called with mu_events held
+    void prune_complete_nolock()
     {
         const auto &ht_it =
             std::remove_if(host_task_events.begin(), host_task_events.end(),
@@ -46,76 +49,78 @@ class SequentialOrder
     }
 
     SequentialOrder(const SequentialOrder &other)
-        : host_task_events(other.host_task_events),
-          submitted_events(other.submitted_events)
     {
-        prune_complete();
+        std::lock_guard<std::mutex> lock(other.mu_events);
+        host_task_events = other.host_task_events;
+        submitted_events = other.submitted_events;
+        prune_complete_nolock();
     }
     SequentialOrder(SequentialOrder &&other)
         : host_task_events{}, submitted_events{}
     {
+        std::lock_guard<std::mutex> lock(other.mu_events);
         host_task_events = std::move(other.host_task_events);
         submitted_events = std::move(other.submitted_events);
-        prune_complete();
+        prune_complete_nolock();
     }
 
     SequentialOrder &operator=(const SequentialOrder &other)
     {
-        host_task_events = other.host_task_events;
-        submitted_events = other.submitted_events;
-
-        prune_complete();
+        if (this != &other) {
+            std::scoped_lock lock(mu_events, other.mu_events);
+            host_task_events = other.host_task_events;
+            submitted_events = other.submitted_events;
+            prune_complete_nolock();
+        }
         return *this;
     }
 
     SequentialOrder &operator=(SequentialOrder &&other)
     {
         if (this != &other) {
+            std::scoped_lock lock(mu_events, other.mu_events);
             host_task_events = std::move(other.host_task_events);
             submitted_events = std::move(other.submitted_events);
-            prune_complete();
+            prune_complete_nolock();
         }
         return *this;
     }
 
     std::size_t get_num_submitted_events() const
     {
+        std::lock_guard<std::mutex> lock(mu_events);
         return submitted_events.size();
     }
 
-    const std::vector<sycl::event> &get_host_task_events()
+    // returns a copy to avoid returning a reference that
+    // could be modified after the lock is released
+    std::vector<sycl::event> get_host_task_events()
     {
-        prune_complete();
+        std::lock_guard<std::mutex> lock(mu_events);
+        prune_complete_nolock();
         return host_task_events;
     }
 
-    /*
-        const std::vector<sycl::event> & get_host_task_events() const {
-            return host_task_events;
-        }
-    */
-
     std::size_t get_num_host_task_events() const
     {
+        std::lock_guard<std::mutex> lock(mu_events);
         return host_task_events.size();
     }
 
-    const std::vector<sycl::event> &get_submitted_events()
+    // returns a copy to avoid returning a reference that
+    // could be modified after the lock is released
+    std::vector<sycl::event> get_submitted_events()
     {
-        prune_complete();
+        std::lock_guard<std::mutex> lock(mu_events);
+        prune_complete_nolock();
         return submitted_events;
     }
 
-    /*
-        const std::vector<sycl::event> & get_submitted_events() const {
-            return submitted_events;
-        }
-    */
-
     void add_to_both_events(const sycl::event &ht_ev,
                             const sycl::event &comp_ev)
     {
-        prune_complete();
+        std::lock_guard<std::mutex> lock(mu_events);
+        prune_complete_nolock();
         if (!is_event_complete(ht_ev))
             host_task_events.push_back(ht_ev);
         if (!is_event_complete(comp_ev))
@@ -125,7 +130,8 @@ class SequentialOrder
     void add_vector_to_both_events(const std::vector<sycl::event> &ht_evs,
                                    const std::vector<sycl::event> &comp_evs)
     {
-        prune_complete();
+        std::lock_guard<std::mutex> lock(mu_events);
+        prune_complete_nolock();
         for (const auto &e : ht_evs) {
             if (!is_event_complete(e))
                 host_task_events.push_back(e);
@@ -138,15 +144,17 @@ class SequentialOrder
 
     void add_to_host_task_events(const sycl::event &ht_ev)
     {
-        prune_complete();
+        std::lock_guard<std::mutex> lock(mu_events);
+        prune_complete_nolock();
         if (!is_event_complete(ht_ev)) {
             host_task_events.push_back(ht_ev);
         }
     }
 
     void add_to_submitted_events(const sycl::event &comp_ev)
     {
-        prune_complete();
+        std::lock_guard<std::mutex> lock(mu_events);
+        prune_complete_nolock();
         if (!is_event_complete(comp_ev)) {
             submitted_events.push_back(comp_ev);
         }
@@ -155,7 +163,8 @@ class SequentialOrder
     template <std::size_t num>
     void add_list_to_host_task_events(const sycl::event (&ht_events)[num])
     {
-        prune_complete();
+        std::lock_guard<std::mutex> lock(mu_events);
+        prune_complete_nolock();
         for (std::size_t i = 0; i < num; ++i) {
             const auto &e = ht_events[i];
             if (!is_event_complete(e))
@@ -166,7 +175,8 @@ class SequentialOrder
     template <std::size_t num>
     void add_list_to_submitted_events(const sycl::event (&comp_events)[num])
     {
-        prune_complete();
+        std::lock_guard<std::mutex> lock(mu_events);
+        prune_complete_nolock();
         for (std::size_t i = 0; i < num; ++i) {
             const auto &e = comp_events[i];
             if (!is_event_complete(e))
@@ -176,8 +186,20 @@ class SequentialOrder
 
     void wait()
     {
-        sycl::event::wait(submitted_events);
-        sycl::event::wait(host_task_events);
-        prune_complete();
+        // snapeshot events outside of mutex to avoid
+        // calling wait inside mutex
+        std::vector<sycl::event> sub_copy;
+        std::vector<sycl::event> ht_copy;
+        {
+            std::lock_guard<std::mutex> lock(mu_events);
+            sub_copy = submitted_events;
+            ht_copy = host_task_events;
+        }
+        sycl::event::wait(sub_copy);
+        sycl::event::wait(ht_copy);
+        {
+            std::lock_guard<std::mutex> lock(mu_events);
+            prune_complete_nolock();
+        }
     }
 };