Hold queues by weak reference in SequentialOrderManager

CHANGELOG.md

@@ -13,6 +13,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 * Added `dpctl::utils::defer_release` to `dpctl4pybind11.hpp`, which calls a release callback once the given events have completed [gh-2359](https://github.com/IntelPython/dpctl/pull/2359)
 * Added `dpctl.SyclQueue.get_last_event` and `dpctl.SyclQueue.is_get_last_event_available`, and the `DPCTLQueue_GetLastEvent` and `DPCTLQueue_GetLastEvent_Available` C-API functions, wrapping the DPC++ `sycl_ext_oneapi_in_order_queue_events` extension [gh-2299](https://github.com/IntelPython/dpctl/pull/2299)
 * Added `dpctl.SyclPlatform.has_aspect_*` properties and the `DPCTLPlatform_HasAspect` C-API function for querying aspects supported by all devices of a platform [gh-2412](https://github.com/IntelPython/dpctl/pull/2412)
+* Added support for weak references to `dpctl.SyclQueue` [gh-2426](https://github.com/IntelPython/dpctl/pull/2426)
 
 ### Changed
 * Object lifetime during offload is now managed by a list of releases deferred until the events that gate them have completed, rather than by a `host_task`. `dpctl` keeps one thread for looking over the list and running the releases that have come due, started by the first release deferred and again in the child of a `fork`, which takes the GIL to do so and lets a release call into the interpreter and free USM allocations [gh-2359](https://github.com/IntelPython/dpctl/pull/2359)
@@ -31,6 +32,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 * Removed stale `select_host_device` and `has_host_device` names from `dpctl.__all__`, which made `from dpctl import *` fail with `AttributeError` [gh-2410](https://github.com/IntelPython/dpctl/pull/2410)
 * A USM allocation that fails now drops the references held for offloaded tasks that have completed, and is attempted once more if that released anything, so that memory only waiting to be given up is not reported as unavailable [gh-2359](https://github.com/IntelPython/dpctl/pull/2359)
 * `dpctl.SyclPlatform.get_devices` no longer prints Unified Runtime warnings for `device_type="custom"` on non-OpenCL platforms, where custom devices are not supported [gh-2306](https://github.com/IntelPython/dpctl/pull/2306)
+* `dpctl.utils.SequentialOrderManager` no longer keeps every queue alive until `clear()`: it holds queues by weak reference and drops the order manager of a queue once the queue is gone and its work is complete, so a program that keeps creating queues no longer runs out of memory [gh-2426](https://github.com/IntelPython/dpctl/pull/2426)
 
 ### Maintenance
 * Removed the `-fno-sycl-id-queries-fit-in-int` compile option, which none of the `dpctl` extensions need since `dpctl.tensor` moved to `dpnp` [gh-2413](https://github.com/IntelPython/dpctl/pull/2413)

docs/doc_sources/api_reference/dpctl/utils.rst

@@ -16,6 +16,14 @@
     Thread-local object mapping each :class:`dpctl.SyclQueue` to an order
     manager, used to ensure sequential ordering of offloaded tasks.
 
+    Queues are held by weak reference, so creating queues does not keep them
+    alive. Equal queues, which share a ``sycl::queue``, share an order manager,
+    which lives while the queue it was looked up by does. Once that queue is
+    gone, the order manager is dropped when the work recorded in it is
+    complete and nothing else holds it; until then, an equal queue is given
+    the same order manager, so that its tasks still follow that work. Nothing
+    waits when a queue is dropped.
+
     Record submitted tasks with ``add_event`` and use ``submitted_events``
     as the dependency list of subsequent submissions. To keep Python objects
     referenced by a task alive until it completes, use

dpctl/_sycl_queue.pxd

@@ -50,6 +50,11 @@ cdef public api class SyclQueue (_SyclQueue) [
 ]:
     """ Python wrapper class for a sycl::queue.
     """
+    # lets dpctl.utils.SequentialOrderManager hold queues by weak reference;
+    # here rather than in _SyclQueue, so that it follows the vtable pointer and
+    # leaves the offsets that extensions built against earlier dpctl use as
+    # they were
+    cdef object __weakref__
     cdef int _init_queue_default(self, int)
     cdef int _init_queue_from__SyclQueue(self, _SyclQueue)
     cdef int _init_queue_from_DPCTLSyclDeviceRef(self, DPCTLSyclDeviceRef, int)

dpctl/_sycl_queue.pyx

@@ -984,6 +984,14 @@ cdef class SyclQueue(_SyclQueue):
                     "be renamed"
                 )
 
+    def __dealloc__(self):
+        # Nothing to release beyond what _SyclQueue does, but defining it
+        # makes Cython give SyclQueue a deallocator of its own, which clears
+        # the weak references to the queue. Cython does not count __weakref__
+        # as a field that needs one, and would use _SyclQueue's instead,
+        # leaving them pointing to freed memory.
+        pass
+
     cdef int _init_queue_from__SyclQueue(self, _SyclQueue other):
         """ Copy data container _SyclQueue fields over.
         """

dpctl/tests/test_sycl_queue.py

@@ -18,6 +18,7 @@
 
 import ctypes
 import sys
+import weakref
 
 import pytest
 
@@ -132,6 +133,23 @@ def test_hashing_of_queue():
     assert queue_dict
 
 
+@pytest.mark.parametrize(
+    "cls", [dpctl.SyclQueue, type("Q", (dpctl.SyclQueue,), {})]
+)
+def test_weak_reference_to_queue(cls):
+    try:
+        q = cls()
+    except dpctl.SyclQueueCreationError:
+        pytest.skip("Default-constructed queue could not be created")
+    gone = []
+    r = weakref.ref(q, gone.append)
+    assert r() is q
+    del q
+    # the queue's deallocator has to clear its weak references
+    assert gone == [r]
+    assert r() is None
+
+
 def test_channeling_device_properties(capsys):
     try:
         q = dpctl.SyclQueue()

dpctl/tests/test_utils.py

@@ -23,6 +23,7 @@
 import threading
 import time
 import warnings
+import weakref
 
 import numpy as np
 import pytest
@@ -659,14 +660,165 @@ def wait(self):
 
 def test_in_order_order_manager_waits_on_queue():
     q = _in_order_queue(cls=_WaitCountingQueue)
-    _som = dpctl.utils.SequentialOrderManager
-    _mngr = _som[q]
+    # the map gives its managers a copy of the queue, so the manager is made
+    # here to see its waits
+    _mngr = _order_manager._InOrderManager(q)
     _mngr.wait()
     assert q.num_waits == 1
     cpy = _mngr.__copy__()
     del cpy
     assert q.num_waits == 2
     del _mngr
+    assert q.num_waits == 3
+
+
+def test_order_manager_clear_waits_on_in_order_queues(monkeypatch):
+    waited = []
+    _wait = _order_manager._InOrderManager.wait
+
+    def wait(self):
+        waited.append(self)
+        return _wait(self)
+
+    monkeypatch.setattr(_order_manager._InOrderManager, "wait", wait)
+    q = _in_order_queue()
+    _som = dpctl.utils.SequentialOrderManager
+    _mngr = _som[q]
     # clear() runs at exit, and must wait on in-order queues as well
     _som.clear()
-    assert q.num_waits > 2
+    assert any(m is _mngr for m in waited)
+
+
+def _thread_map():
+    return dpctl.utils.SequentialOrderManager._get_thread_map()
+
+
+@pytest.mark.parametrize("in_order", [False, True], ids=["ooo", "in_order"])
+def test_order_manager_does_not_keep_queues_alive(in_order):
+    _som = dpctl.utils.SequentialOrderManager
+    _som.clear()
+    refs = []
+    for _ in range(16):
+        q = _in_order_queue() if in_order else _queue()
+        _mngr = _som[q]
+        buf = dpctl.memory.MemoryUSMDevice(1024, queue=q)
+        _mngr.add_event(q.memset_async(buf, 1, buf.nbytes))
+        # holds nothing: buf would keep q alive until dpctl's thread drops it
+        _mngr.add_cleanup_event(q.keep_args_alive(()))
+        _mngr.wait()
+        refs.append(weakref.ref(q))
+        del q, _mngr, buf
+        # each lookup drops the managers of the queues gone before it
+        assert len(_thread_map().managers) <= 2
+    assert all(r() is None for r in refs)
+    m = _thread_map()
+    assert len(m.dead) == 1
+    q = _in_order_queue() if in_order else _queue()
+    _som[q]
+    assert len(m.managers) == 1 and len(m.owned) == 1
+    assert not m.retired and not m.dead and not m.dropped
+    _som.clear()
+
+
+def test_order_manager_hands_pending_work_to_equal_queue():
+    _som = dpctl.utils.SequentialOrderManager
+    _som.clear()
+    q1 = _queue()
+    # the same sycl::queue in another object
+    q2 = dpctl.SyclQueue(q1._get_capsule())
+    assert q1 == q2 and q1 is not q2
+    _mngr = _som[q1]
+    assert _som[q2] is _mngr
+    # as if the work recorded through q1 were not complete when q1 is gone
+    _mngr._is_idle = lambda: False
+    _mngr_ref = weakref.ref(_mngr)
+    del q1, _mngr
+    _som[_queue()]
+    # the map holds it, as its work is not complete
+    assert _mngr_ref() is not None
+    assert _som[q2] is _mngr_ref()
+    _som.clear()
+
+
+def test_order_manager_drops_idle_manager_of_queue_gone():
+    _som = dpctl.utils.SequentialOrderManager
+    _som.clear()
+    q1 = _queue()
+    q2 = dpctl.SyclQueue(q1._get_capsule())
+    _mngr = _som[q1]
+    _mngr.wait()
+    _mngr_ref = weakref.ref(_mngr)
+    del q1, _mngr
+    # nothing holds it once the map drops it
+    _som[q2]
+    assert _mngr_ref() is None
+    m = _thread_map()
+    assert len(m.owned) == 1 and not m.retired and not m.dropped
+    _som.clear()
+
+
+def test_order_manager_keeps_manager_held_for_equal_queue():
+    _som = dpctl.utils.SequentialOrderManager
+    _som.clear()
+    q1 = _queue()
+    q2 = dpctl.SyclQueue(q1._get_capsule())
+    _som[q1]
+    # looked up by q2, but the map has it under q1
+    _mngr = _som[q2]
+    _mngr.wait()
+    del q1
+    # the map drops it, its work being complete, while it is still in use
+    _som[_queue()]
+    assert _mngr._dropped
+    # work recorded in it now has to be ordered before the work on q2
+    assert _som[q2] is _mngr
+    assert not _mngr._dropped
+    _som.clear()
+
+
+def test_order_manager_orders_work_of_queue_gone():
+    _som = dpctl.utils.SequentialOrderManager
+    _som.clear()
+    q2 = _queue()
+    buf = dpctl.memory.MemoryUSMDevice(64 * 1024 * 1024, queue=q2)
+    host = np.empty(buf.nbytes, dtype=np.uint8)
+    for _ in range(5):
+        q2.memset(buf, 0, buf.nbytes)
+        _som.clear()
+        # the same sycl::queue in another object, which is the key
+        q1 = dpctl.SyclQueue(q2._get_capsule())
+        _mngr = _som[q1]
+        _fill_in_steps(q1, buf, 8, _mngr)
+        del _mngr
+        q1_ref = weakref.ref(q1)
+        del q1
+        assert q1_ref() is None
+        _mngr = _som[q2]
+        ev = q2.memset_async(buf, 9, buf.nbytes, _mngr.submitted_events)
+        _mngr.add_event(ev)
+        q2.memcpy_async(host, buf, buf.nbytes, [ev]).wait()
+        assert np.all(host == 9)
+        del _mngr
+    _som.clear()
+
+
+def test_order_manager_queue_gone_on_other_thread(monkeypatch):
+    _som = dpctl.utils.SequentialOrderManager
+    _som.clear()
+    q = _queue()
+    _mngr = _som[q]
+    monkeypatch.setattr(_mngr, "_is_idle", lambda: False)
+    holder = [q]
+    del q
+    # the last reference is dropped on another thread, which neither waits
+    # nor changes this thread's map
+    t = threading.Thread(target=holder.clear)
+    t.start()
+    t.join(timeout=60)
+    assert not t.is_alive()
+    m = _thread_map()
+    assert len(m.dead) == 1 and len(m.managers) == 1
+    _som[_queue()]
+    assert [om for om in m.retired.values()] == [_mngr]
+    _som.clear()
+    assert not m.retired