Migrate host_tasks to a thread pool based approach

Inspired by comment in: AdaptiveCpp/AdaptiveCpp#1915

we move away from use of host_task, which enables compatibility with AdaptiveCpp

Author: ndgrigorian

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/_host_task_util.hpp dpctl/_async_dec_ref.hppdpctl/_host_task_util.hpp renamed to 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,9 +33,12 @@
 #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"
 
@@ -49,16 +52,21 @@ DPCTLSyclEventRef async_dec_ref(DPCTLSyclQueueRef QRef,
     using dpctl::syclinterface::unwrap;
     using dpctl::syclinterface::wrap;
 
-    sycl::queue *q = unwrap<sycl::queue>(QRef);
+    // `QRef` is kept in the signature for API compatibility
+    (void)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,12 +83,12 @@ 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 a dummy event for API compatibility
+        auto e_ptr = new sycl::event();
         return wrap<sycl::event>(e_ptr);
     } catch (const std::exception &e) {
         static constexpr int result_std_exception = 1;

dpctl/_sycl_queue.pyx

@@ -92,7 +92,7 @@ import collections.abc
 import logging
 
 
-cdef extern from "_host_task_util.hpp":
+cdef extern from "_async_dec_ref.hpp":
     DPCTLSyclEventRef async_dec_ref(
         DPCTLSyclQueueRef, PyObject **,
         size_t, DPCTLSyclEventRef *, size_t, int *
@@ -1184,16 +1184,18 @@ 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.
         """
         cdef size_t nDE = len(dEvents)
@@ -1236,7 +1238,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 cleanup")
 
         return SyclEvent._create(htERef)
 

dpctl/apis/include/detail/keep_alive_pool.hpp

@@ -0,0 +1,125 @@
+//===--- keep_alive_pool.hpp - keeps owners alive during offload ----------===//
+//
+//                      Data Parallel Control (dpctl)
+//
+// Copyright 2026 Intel Corporation
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//    http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+//
+//===----------------------------------------------------------------------===//
+///
+/// \file
+/// A fixed-size pool of threads that wait on SYCL events and then run a
+/// callable for maintaining Python object lifetime during offloaded tasks.
+///
+//===----------------------------------------------------------------------===//
+
+#pragma once
+
+#include <condition_variable>
+#include <cstddef>
+#include <exception>
+#include <functional>
+#include <mutex>
+#include <queue>
+#include <thread>
+#include <utility>
+#include <vector>
+
+#include <sycl/sycl.hpp>
+
+namespace dpctl
+{
+namespace detail
+{
+
+class KeepAlivePool
+{
+public:
+    /*!
+     * @brief Number of waiter threads.
+     */
+    static constexpr std::size_t num_threads = 4;
+
+    static KeepAlivePool &get()
+    {
+        // deliberately leaked: workers are detached and hold a bare `this`, so
+        // the pool must outlive them
+        static KeepAlivePool *instance = new KeepAlivePool();
+        return *instance;
+    }
+
+    /*!
+     * @brief Run `task` once every event in `depends` has completed.
+     *
+     * `task` must own everything it releases -- move USM `shared_ptr` copies,
+     * `sycl::buffer` handles, or `PyObject *` references into it. It runs on a
+     * pool thread without the GIL, so it must take the GIL itself if it touches
+     * Python.
+     */
+    void submit(std::vector<sycl::event> depends, std::function<void()> task)
+    {
+        {
+            std::lock_guard<std::mutex> lock(queue_mutex_);
+            tasks_.emplace(std::move(depends), std::move(task));
+        }
+        condition_.notify_one();
+    }
+
+    KeepAlivePool(const KeepAlivePool &) = delete;
+    KeepAlivePool &operator=(const KeepAlivePool &) = delete;
+    ~KeepAlivePool() = delete;
+
+private:
+    KeepAlivePool()
+    {
+        for (std::size_t i = 0; i < num_threads; ++i) {
+            std::thread(&KeepAlivePool::run, this).detach();
+        }
+    }
+
+    void run()
+    {
+        for (;;) {
+            std::pair<std::vector<sycl::event>, std::function<void()>> item;
+            {
+                std::unique_lock<std::mutex> lock(queue_mutex_);
+                condition_.wait(lock, [this] { return !tasks_.empty(); });
+                item = std::move(tasks_.front());
+                tasks_.pop();
+            }
+
+            try {
+                sycl::event::wait(item.first);
+            } catch (const std::exception &) {
+                // run the task anyway: an async error must not strand the
+                // task or else it may leak
+            }
+
+            try {
+                item.second();
+            } catch (const std::exception &) {
+                // a throwing task must not take down the worker or later
+                // tasks will be lost
+            }
+        }
+    }
+
+    std::queue<std::pair<std::vector<sycl::event>, std::function<void()>>>
+        tasks_;
+    std::mutex queue_mutex_;
+    std::condition_variable condition_;
+};
+
+} // namespace detail
+} // namespace dpctl

dpctl/apis/include/dpctl4pybind11.hpp

@@ -25,6 +25,7 @@
 
 #pragma once
 
+#include "detail/keep_alive_pool.hpp"
 #include "dpctl_capi.h"
 
 #include <atomic>
@@ -823,46 +824,26 @@ sycl::event keep_args_alive(sycl::queue &q,
         }
     }
 
-    bool use_depends = true;
-    sycl::event host_task_ev;
-
-    if (n_usm_owners_held > 0) {
-        host_task_ev = q.submit([&](sycl::handler &cgh) {
-            if (use_depends) {
-                cgh.depends_on(depends);
-                use_depends = false;
-            }
-            else {
-                cgh.depends_on(host_task_ev);
-            }
-            cgh.host_task([shp_usm = std::move(shp_usm)]() {
-                // no body, but shared pointers are captured in
-                // the lambda, ensuring that USM allocation is
-                // kept alive
-            });
-        });
-    }
+    if (n_usm_owners_held > 0 || n_objects_held > 0) {
+        dpctl::detail::KeepAlivePool::get().submit(
+            depends, [n_usm_owners_held, shp_usm = std::move(shp_usm),
+                      n_objects_held, shp_arr = std::move(shp_arr)]() mutable {
+                for (std::size_t i = 0; i < n_usm_owners_held; ++i) {
+                    shp_usm[i].reset();
+                }
 
-    if (n_objects_held > 0) {
-        host_task_ev = q.submit([&](sycl::handler &cgh) {
-            if (use_depends) {
-                cgh.depends_on(depends);
-                use_depends = false;
-            }
-            else {
-                cgh.depends_on(host_task_ev);
-            }
-            cgh.host_task([n_objects_held, shp_arr = std::move(shp_arr)]() {
-                py::gil_scoped_acquire acquire;
+                if (n_objects_held > 0) {
+                    py::gil_scoped_acquire acquire;
 
-                for (std::size_t i = 0; i < n_objects_held; ++i) {
-                    shp_arr[i]->dec_ref();
+                    for (std::size_t i = 0; i < n_objects_held; ++i) {
+                        shp_arr[i]->dec_ref();
+                    }
                 }
             });
-        });
     }
 
-    return host_task_ev;
+    // return dummy event for API compatibility
+    return sycl::event{};
 }
 
 /*! @brief Check if all allocation queues are the same as the