Migrate host_task object management to a thread pool based approach

CHANGELOG.md

@@ -11,9 +11,15 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 * Added a number of `sycl::device` info queries to `dpctl.SyclDevice` [gh-2324](https://github.com/IntelPython/dpctl/pull/2324)
 * Added `sycl::info::context` queries `sycl_platform`, `atomic_memory_order_capabilities`, `atomic_fence_order_capabilities`, `atomic_memory_scope_capabilities`, and `atomic_fence_scope_capabilities` to `dpctl.SyclContext` [gh-2354](https://github.com/IntelPython/dpctl/pull/2354)
 * Added `create_kernel_bundle_from_sycl_source`, `is_sycl_source_compilation_available`, and `dpctl.SyclDevice.can_compile` for supporting the creation of `dpctl.SyclKernelBundle`s from SYCL source strings via DPC++ extension, as well as corresponding C-API functions to support it [gh-2206](https://github.com/IntelPython/dpctl/pull/2206)
+* Added `dpctl.keep_args_alive` free function, and `add_event` method to the order manager
+
+### Deprecated
+* Deprecated `dpctl.SyclQueue._submit_keep_args_alive` in favor of `dpctl.keep_args_alive` [gh-2359](https://github.com/IntelPython/dpctl/pull/2359)
+* Deprecated the order manager's `add_event_pair`, `host_task_events` and `num_host_task_events`, as `host_task` is no longer used for managing object lifetimes [gh-2359](https://github.com/IntelPython/dpctl/pull/2359)
 
 ### Changed
 * Bump minimum NumPy version to 1.26 [gh-2192](https://github.com/IntelPython/dpctl/pull/2192)
+* Implemented a thread pool to manage object lifetime during offload rather than use `host_task` [gh-2359](https://github.com/IntelPython/dpctl/pull/2359)
 * Rewrote USM Python examples into a single example [gh-2292](https://github.com/IntelPython/dpctl/pull/2292)
 * Registered `DPCTL_PARTITION_AFFINITY_DOMAIN_UNKNOWN` enumerator when `DPCTLDevice_GetPartitionAffinityDomains` receives an unrecognized value from the SYCL runtime [gh-2324](https://github.com/IntelPython/dpctl/pull/2324)
 

docs/doc_sources/api_reference/dpctl/index.rst

@@ -86,6 +86,14 @@
     SyclQueueCreationError
     SyclSubDeviceCreationError
 
+.. rubric:: Lifetime management
+
+.. autosummary::
+    :toctree: generated
+    :nosignatures:
+
+    keep_args_alive
+
 .. rubric:: Utilities
 
 .. autosummary::

docs/doc_sources/api_reference/dpctl/utils.rst

@@ -15,3 +15,13 @@
 
     Thread-local object mapping each :class:`dpctl.SyclQueue` to an order
     manager, used to ensure sequential ordering of offloaded tasks.
+
+    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
+    :func:`dpctl.keep_args_alive`.
+
+    .. deprecated::
+        ``add_event_pair``, ``host_task_events`` and ``num_host_task_events``
+        are deprecated. Tasks are no longer paired with a host task event,
+        so ``add_event`` takes the computational event alone.

dpctl/CMakeLists.txt

@@ -199,8 +199,12 @@ endforeach()
 set(_cy_file ${CMAKE_CURRENT_SOURCE_DIR}/_sycl_queue.pyx)
 get_filename_component(_trgt ${_cy_file} NAME_WLE)
 build_dpctl_ext(${_trgt} ${_cy_file} "dpctl" SYCL)
-# _sycl_queue include _host_task_util.hpp
-target_include_directories(${_trgt} PRIVATE ${CMAKE_CURRENT_SOURCE_DIR})
+# _sycl_queue includes _async_dec_ref.hpp, which includes
+# detail/keep_alive_pool.hpp from the public include directory
+target_include_directories(${_trgt} PRIVATE
+    ${CMAKE_CURRENT_SOURCE_DIR}
+    ${CMAKE_CURRENT_SOURCE_DIR}/apis/include
+)
 target_link_libraries(DpctlCAPI INTERFACE ${_trgt}_headers)
 
 add_subdirectory(program)

dpctl/__init__.py

@@ -57,6 +57,7 @@
     SyclQueue,
     SyclQueueCreationError,
     WorkGroupMemory,
+    keep_args_alive,
 )
 from ._sycl_queue_manager import get_device_cached_queue
 from ._sycl_timer import SyclTimer
@@ -114,6 +115,7 @@
     "WorkGroupMemory",
     "LocalAccessor",
     "RawKernelArg",
+    "keep_args_alive",
 ]
 __all__ += [
     "get_device_cached_queue",

dpctl/_host_task_util.hpp → dpctl/_async_dec_ref.hpp

@@ -1,4 +1,4 @@
-//===--- _host_tasl_util.hpp - Implements async DECREF =//
+//===--- _async_dec_ref.hpp - Implements async DECREF ---------------------===//
 //
 //                      Data Parallel Control (dpctl)
 //
@@ -19,11 +19,11 @@
 //===----------------------------------------------------------------------===//
 ///
 /// \file
-/// This file implements a utility function to schedule host task to a sycl
-/// queue depending on given array of sycl events to decrement reference counts
-/// for the given array of Python objects.
+/// This file implements a utility function to decrement reference counts for a
+/// given array of Python objects once a given array of sycl events has
+/// completed.
 ///
-/// N.B.: The host task attempts to acquire GIL, so queue wait, event wait and
+/// N.B.: The deferred work acquires the GIL, so queue wait, event wait and
 /// other synchronization mechanisms should be called after releasing the GIL to
 /// avoid deadlocks.
 ///
@@ -33,32 +33,40 @@
 #include <exception>
 #include <stddef.h>
 #include <sycl/sycl.hpp>
+#include <utility>
+#include <vector>
 
 #include "Python.h"
 
+#include "detail/keep_alive_pool.hpp"
 #include "syclinterface/dpctl_data_types.h"
 #include "syclinterface/dpctl_sycl_type_casters.hpp"
 
-DPCTLSyclEventRef async_dec_ref(DPCTLSyclQueueRef QRef,
-                                PyObject **obj_array,
-                                size_t obj_array_size,
-                                DPCTLSyclEventRef *depERefs,
-                                size_t nDepERefs,
-                                int *status)
+/*!
+ * @brief Schedule DECREFs of `obj_array` for once `depERefs` have completed.
+ *
+ * Sets `*status` to 0 on success and 1 if scheduling threw.
+ */
+void async_dec_ref(PyObject **obj_array,
+                   size_t obj_array_size,
+                   DPCTLSyclEventRef *depERefs,
+                   size_t nDepERefs,
+                   int *status)
 {
     using dpctl::syclinterface::unwrap;
-    using dpctl::syclinterface::wrap;
-
-    sycl::queue *q = unwrap<sycl::queue>(QRef);
 
     std::vector<PyObject *> obj_vec(obj_array, obj_array + obj_array_size);
 
     try {
-        sycl::event ht_ev = q->submit([&](sycl::handler &cgh) {
-            for (size_t ev_id = 0; ev_id < nDepERefs; ++ev_id) {
-                cgh.depends_on(*(unwrap<sycl::event>(depERefs[ev_id])));
-            }
-            cgh.host_task([obj_array_size, obj_vec]() {
+        std::vector<sycl::event> depends;
+        depends.reserve(nDepERefs);
+        for (size_t ev_id = 0; ev_id < nDepERefs; ++ev_id) {
+            depends.push_back(*(unwrap<sycl::event>(depERefs[ev_id])));
+        }
+
+        dpctl::detail::KeepAlivePool::get().submit(
+            std::move(depends),
+            [obj_array_size, obj_vec = std::move(obj_vec)]() {
                 const bool initialized = Py_IsInitialized();
 #if PY_VERSION_HEX < 0x30d0000
                 const bool finalizing = _Py_IsFinalizing();
@@ -75,22 +83,38 @@ DPCTLSyclEventRef async_dec_ref(DPCTLSyclQueueRef QRef,
                     PyGILState_Release(gstate);
                 }
             });
-        });
 
         static constexpr int result_ok = 0;
-
         *status = result_ok;
-        auto e_ptr = new sycl::event(ht_ev);
-        return wrap<sycl::event>(e_ptr);
     } catch (const std::exception &e) {
         static constexpr int result_std_exception = 1;
-
         *status = result_std_exception;
-        return nullptr;
     }
+}
 
-    static constexpr int result_other_abnormal = 2;
+/*!
+ * @brief Event-returning form of `async_dec_ref`.
+ *
+ * Returns a default-constructed event.
+ * Returns nullptr on failure, with `*status` set.
+ */
+DPCTLSyclEventRef async_dec_ref_event(DPCTLSyclQueueRef QRef,
+                                      PyObject **obj_array,
+                                      size_t obj_array_size,
+                                      DPCTLSyclEventRef *depERefs,
+                                      size_t nDepERefs,
+                                      int *status)
+{
+    using dpctl::syclinterface::wrap;
+
+    (void)QRef;
+
+    async_dec_ref(obj_array, obj_array_size, depERefs, nDepERefs, status);
+
+    if (*status != 0) {
+        return nullptr;
+    }
 
-    *status = result_other_abnormal;
-    return nullptr;
+    auto e_ptr = new sycl::event();
+    return wrap<sycl::event>(e_ptr);
 }

dpctl/_sycl_queue.pyx

@@ -85,21 +85,27 @@ from cpython.buffer cimport (
     PyObject_CheckBuffer,
     PyObject_GetBuffer,
 )
-from cpython.ref cimport Py_INCREF, PyObject
+from cpython.ref cimport Py_DECREF, Py_INCREF, PyObject
 from libc.stdlib cimport free, malloc
 
 import collections.abc
 import logging
+import warnings
 
 
-cdef extern from "_host_task_util.hpp":
-    DPCTLSyclEventRef async_dec_ref(
+cdef extern from "_async_dec_ref.hpp":
+    void async_dec_ref(
+        PyObject **, size_t, DPCTLSyclEventRef *, size_t, int *
+    ) nogil
+    # deprecated, retained for the queue-bound _submit_keep_args_alive
+    DPCTLSyclEventRef async_dec_ref_event(
         DPCTLSyclQueueRef, PyObject **,
         size_t, DPCTLSyclEventRef *, size_t, int *
     ) nogil
 
 
 __all__ = [
+    "keep_args_alive",
     "SyclQueue",
     "SyclKernelInvalidRangeError",
     "SyclKernelSubmitError",
@@ -1174,6 +1180,9 @@ cdef class SyclQueue(_SyclQueue):
         Keeps objects in ``args`` alive until tasks associated with events
         complete.
 
+        Deprecated since dpctl 0.23.0. Use :func:`dpctl.keep_args_alive`
+        instead, which is not bound to a queue and returns nothing.
+
         Args:
             args(object):
                 Python object to keep alive.
@@ -1184,18 +1193,28 @@ cdef class SyclQueue(_SyclQueue):
                 working on Python objects collected in ``args``.
         Returns:
             dpctl.SyclEvent
-               The event associated with the submission of host task.
+               An already-complete event. No task is submitted to the queue,
+               so there is nothing to wait for here; ``events`` are what say
+               when the work reading ``args`` is done.
 
-        Increments reference count of ``args`` and schedules asynchronous
-        ``host_task`` to decrement the count once dependent events are
+        Increments reference count of ``args`` and schedules the matching
+        decrement to run on a background thread once dependent events are
         complete.
 
         .. note::
-            The ``host_task`` attempts to acquire Python GIL, and it is
-            known to be unsafe during interpreter shutdown sequence. It is
-            thus strongly advised to ensure that all submitted ``host_task``
+            The deferred decrement attempts to acquire Python GIL, which is
+            known to be unsafe during the interpreter shutdown sequence. It
+            is thus strongly advised to ensure that all dependent events
             complete before the end of the Python script.
         """
+        warnings.warn(
+            "dpctl.SyclQueue._submit_keep_args_alive is deprecated and will "
+            "be removed in a future release. Use dpctl.keep_args_alive "
+            "instead, which is not bound to a queue and returns nothing.",
+            DeprecationWarning,
+            stacklevel=2,
+        )
+
         cdef size_t nDE = len(dEvents)
         cdef DPCTLSyclEventRef *depEvents = NULL
         cdef PyObject *args_raw = NULL
@@ -1225,7 +1244,7 @@ cdef class SyclQueue(_SyclQueue):
         # schedule decrement
         args_raw = <PyObject *>args
 
-        htERef = async_dec_ref(
+        htERef = async_dec_ref_event(
             self.get_queue_ref(),
             &args_raw, 1,
             depEvents, nDE, &status
@@ -1236,7 +1255,7 @@ cdef class SyclQueue(_SyclQueue):
             with nogil:
                 DPCTLEvent_Wait(htERef)
             DPCTLEvent_Delete(htERef)
-            raise RuntimeError("Could not submit keep_args_alive host_task")
+            raise RuntimeError("Could not schedule keep_args_alive")
 
         return SyclEvent._create(htERef)
 
@@ -1278,7 +1297,7 @@ cdef class SyclQueue(_SyclQueue):
             as unified address space pointers.
 
             One way of accomplishing this is to use
-            :meth:`dpctl.SyclQueue._submit_keep_args_alive`.
+            :func:`dpctl.keep_args_alive`.
         """
         cdef void **kargs = NULL
         cdef _arg_data_type *kargty = NULL
@@ -2024,3 +2043,79 @@ cdef class RawKernelArg:
         as a ``size_t``.
         """
         return <size_t>self._arg_ref
+
+
+def keep_args_alive(args, depends):
+    """keep_args_alive(args, depends)
+
+    Keep objects in ``args`` alive until the tasks associated with
+    ``depends`` complete.
+
+    Args:
+        args (object):
+            Python object to keep alive, typically a tuple of the arguments
+            passed to an offloaded task.
+        depends (List[dpctl.SyclEvent]):
+            Gating events. The objects in ``args`` are released once every
+            event in ``depends`` has completed.
+
+    Returns:
+        None
+
+    Increments the reference count of ``args`` and schedules the matching
+    decrement to run on a background thread once every event in ``depends``
+    is complete. The reference is guaranteed to be held for the whole span
+    in between, so the objects cannot be collected while offloaded tasks are
+    still reading them.
+
+    This function is not bound to a queue: the gating events fully determine
+    when the objects may be released.
+
+    :Example:
+        .. code-block:: python
+
+            import dpctl
+
+            q = dpctl.SyclQueue()
+            e = q.submit_async(kernel, [x_usm], [n])
+            dpctl.keep_args_alive((x_usm,), [e])
+
+    .. note::
+        The deferred decrement attempts to acquire the Python GIL, which is
+        known to be unsafe during the interpreter shutdown sequence. It is
+        thus strongly advised to ensure that all events in ``depends``
+        complete before the end of the Python script.
+    """
+    cdef size_t nDE = len(depends)
+    cdef DPCTLSyclEventRef *depEvents = NULL
+    cdef PyObject *args_raw = NULL
+    cdef int status = -1
+
+    if nDE > 0:
+        depEvents = (
+            <DPCTLSyclEventRef*>malloc(nDE*sizeof(DPCTLSyclEventRef))
+        )
+        if not depEvents:
+            raise MemoryError()
+        for idx, de in enumerate(depends):
+            if isinstance(de, SyclEvent):
+                depEvents[idx] = (<SyclEvent>de).get_event_ref()
+            else:
+                free(depEvents)
+                raise TypeError(
+                    "A sequence of dpctl.SyclEvent is expected"
+                )
+
+    # increment reference counts to list of arguments
+    Py_INCREF(args)
+    args_raw = <PyObject *>args
+
+    # schedule decrement
+    async_dec_ref(&args_raw, 1, depEvents, nDE, &status)
+
+    free(depEvents)
+    if status != 0:
+        # the deferred decrement was never scheduled, so undo the increment
+        # here rather than leak the reference
+        Py_DECREF(args)
+        raise RuntimeError("Could not schedule keep_args_alive")

dpctl/_sycl_timer.py

@@ -78,7 +78,7 @@ def get_event(self):
         ev = self._submit_empty_task_fn(
             sycl_queue=self.queue, depends=self._order_manager.submitted_events
         )
-        self._order_manager.add_event_pair(ev, ev)
+        self._order_manager.add_event(ev)
         return ev