Migrate host_task object management to a thread pool based approach

CHANGELOG.md

@@ -11,9 +11,16 @@ 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 [gh-2359](https://github.com/IntelPython/dpctl/pull/2359)
+* Added `DPCTL_KEEP_ALIVE_POOL_SIZE` CMake option for setting the number of threads that keep objects alive during offload [gh-2359](https://github.com/IntelPython/dpctl/pull/2359)
+
+### 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)
 

CMakeLists.txt

@@ -48,6 +48,19 @@ option(
 size of shared object with offloading sections"
      OFF
 )
+set(DPCTL_KEEP_ALIVE_POOL_SIZE
+    "4"
+    CACHE STRING
+    "Number of threads in the pool that keeps Python objects used by \
+offloaded tasks alive until the tasks complete"
+)
+
+if (NOT DPCTL_KEEP_ALIVE_POOL_SIZE MATCHES "^[1-9][0-9]*$")
+    message(FATAL_ERROR
+        "Invalid value for DPCTL_KEEP_ALIVE_POOL_SIZE: \
+\"${DPCTL_KEEP_ALIVE_POOL_SIZE}\". Expected a positive integer."
+    )
+endif()
 
 find_package(IntelSYCL REQUIRED PATHS ${CMAKE_SOURCE_DIR}/cmake NO_DEFAULT_PATH)
 

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:: 0.23.0
+        ``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.

docs/doc_sources/beginners_guides/installation.rst

@@ -247,6 +247,24 @@ devices at the same time:
 
     python scripts/build_locally.py --verbose --target-cuda --target-hip=gfx1030
 
+Configuring the object management thread pool
+---------------------------------------------
+
+To keep Python objects used by an offloaded task alive until tasks complete,
+:py:mod:`dpctl` maintains a pool of threads that wait on the task's events.
+The pool uses four threads by default. The size is fixed when ``dpctl`` is
+compiled and can be changed with the ``--keep-alive-pool-size`` argument:
+
+.. code-block:: bash
+
+    python scripts/build_locally.py --verbose --keep-alive-pool-size=8
+
+Alternatively, you use the ``DPCTL_KEEP_ALIVE_POOL_SIZE`` CMake option:
+
+.. code-block:: bash
+
+    python scripts/build_locally.py --verbose --cmake-opts="-DDPCTL_KEEP_ALIVE_POOL_SIZE=8"
+
 Running Examples and Tests
 ==========================
 

dpctl/CMakeLists.txt

@@ -199,8 +199,16 @@ 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
+)
+# _sycl_queue owns the pool, so only it needs to know the pool size
+target_compile_definitions(${_trgt} PRIVATE
+    DPCTL_KEEP_ALIVE_POOL_SIZE=${DPCTL_KEEP_ALIVE_POOL_SIZE}
+)
 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/_async_dec_ref.hpp

@@ -0,0 +1,134 @@
+//===--- _async_dec_ref.hpp - Implements async DECREF ---------------------===//
+//
+//                      Data Parallel Control (dpctl)
+//
+// Copyright 2022 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
+/// 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 deferred work acquires the GIL, so queue wait, event wait and
+/// other synchronization mechanisms should be called after releasing the GIL to
+/// avoid deadlocks.
+///
+//===----------------------------------------------------------------------===//
+
+#pragma once
+#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"
+
+/*!
+ * @brief Address of the `KeepAlivePool`.
+ *
+ * Returns nullptr if the pool could not be created.
+ */
+void *keep_alive_pool_ptr()
+{
+    try {
+        return static_cast<void *>(
+            &dpctl::detail::KeepAlivePool::local_instance());
+    } catch (const std::exception &e) {
+        return nullptr;
+    }
+}
+
+/*!
+ * @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;
+
+    std::vector<PyObject *> obj_vec(obj_array, obj_array + obj_array_size);
+
+    try {
+        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::local_instance().submit(
+            std::move(depends),
+            [obj_array_size, obj_vec = std::move(obj_vec)]() {
+#if PY_VERSION_HEX < 0x30d0000
+                const bool finalizing = _Py_IsFinalizing();
+#else
+                const bool finalizing = Py_IsFinalizing();
+#endif
+                // if the main thread has not finalized the interpreter yet
+                if (!finalizing) {
+                    PyGILState_STATE gstate;
+                    gstate = PyGILState_Ensure();
+                    for (size_t i = 0; i < obj_array_size; ++i) {
+                        Py_DECREF(obj_vec[i]);
+                    }
+                    PyGILState_Release(gstate);
+                }
+            });
+
+        static constexpr int result_ok = 0;
+        *status = result_ok;
+    } catch (const std::exception &e) {
+        static constexpr int result_std_exception = 1;
+        *status = result_std_exception;
+    }
+}
+
+/*!
+ * @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;
+    }
+
+    auto e_ptr = new sycl::event();
+    return wrap<sycl::event>(e_ptr);
+}