Apply review feedback

Also change pool getter to avoid any dpctl4pybind11 including extensions from having their own thread pools

Author: ndgrigorian

CHANGELOG.md

@@ -11,7 +11,8 @@ 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
+* 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)

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/utils.rst

@@ -21,7 +21,7 @@
     referenced by a task alive until it completes, use
     :func:`dpctl.keep_args_alive`.
 
-    .. deprecated::
+    .. 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

@@ -205,6 +205,10 @@ 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/_async_dec_ref.hpp

@@ -42,6 +42,21 @@
 #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.
  *
@@ -64,17 +79,16 @@ void async_dec_ref(PyObject **obj_array,
             depends.push_back(*(unwrap<sycl::event>(depERefs[ev_id])));
         }
 
-        dpctl::detail::KeepAlivePool::get().submit(
+        dpctl::detail::KeepAlivePool::local_instance().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();
 #else
                 const bool finalizing = Py_IsFinalizing();
 #endif
                 // if the main thread has not finalized the interpreter yet
-                if (initialized && !finalizing) {
+                if (!finalizing) {
                     PyGILState_STATE gstate;
                     gstate = PyGILState_Ensure();
                     for (size_t i = 0; i < obj_array_size; ++i) {

dpctl/_sycl_queue.pyx

@@ -102,6 +102,7 @@ cdef extern from "_async_dec_ref.hpp":
         DPCTLSyclQueueRef, PyObject **,
         size_t, DPCTLSyclEventRef *, size_t, int *
     ) nogil
+    void *keep_alive_pool_ptr() nogil
 
 
 __all__ = [
@@ -1851,6 +1852,10 @@ cdef api SyclQueue SyclQueue_Make(DPCTLSyclQueueRef QRef):
     cdef DPCTLSyclQueueRef copied_QRef = DPCTLQueue_Copy(QRef)
     return SyclQueue._create(copied_QRef)
 
+
+cdef api void *KeepAlivePool_Get() noexcept nogil:
+    return keep_alive_pool_ptr()
+
 cdef class _WorkGroupMemory:
     def __dealloc__(self):
         if(self._mem_ref):

dpctl/apis/include/detail/keep_alive_pool.hpp

@@ -38,6 +38,10 @@
 
 #include <sycl/sycl.hpp>
 
+#ifndef DPCTL_KEEP_ALIVE_POOL_SIZE
+#define DPCTL_KEEP_ALIVE_POOL_SIZE 4
+#endif
+
 namespace dpctl
 {
 namespace detail
@@ -49,9 +53,15 @@ class KeepAlivePool
     /*!
      * @brief Number of waiter threads.
      */
-    static constexpr std::size_t num_threads = 4;
+    static constexpr std::size_t num_threads = DPCTL_KEEP_ALIVE_POOL_SIZE;
+
+    static_assert(num_threads > 0,
+                  "DPCTL_KEEP_ALIVE_POOL_SIZE must be greater than zero");
 
-    static KeepAlivePool &get()
+    /*!
+     * @brief The instance belonging to this shared object.
+     */
+    static KeepAlivePool &local_instance()
     {
         // deliberately leaked: workers are detached and hold a bare `this`, so
         // the pool must outlive them

dpctl/apis/include/dpctl4pybind11.hpp

@@ -47,6 +47,21 @@ namespace dpctl
 namespace detail
 {
 
+/*!
+ * @brief Whether the interpreter can still be called into.
+ *
+ * Acquiring the GIL once finalization has begun does not return, so work
+ * deferred to a thread must check this before touching Python.
+ */
+inline bool interpreter_is_live()
+{
+#if PY_VERSION_HEX < 0x30d0000
+    return !_Py_IsFinalizing();
+#else
+    return !Py_IsFinalizing();
+#endif
+}
+
 class dpctl_capi
 {
 public:
@@ -174,15 +189,7 @@ class dpctl_capi
     {
         void operator()(py::object *p) const
         {
-            const bool initialized = Py_IsInitialized();
-#if PY_VERSION_HEX < 0x30d0000
-            const bool finalizing = _Py_IsFinalizing();
-#else
-            const bool finalizing = Py_IsFinalizing();
-#endif
-            const bool guard = initialized && !finalizing;
-
-            if (guard) {
+            if (interpreter_is_live()) {
                 delete p;
             }
         }
@@ -295,6 +302,25 @@ class dpctl_capi
     dpctl_capi &operator=(dpctl_capi &&) = default;
 
 }; // struct dpctl_capi
+
+/*!
+ * @brief The `KeepAlivePool` singleton, owned by `dpctl._sycl_queue`.
+ */
+inline KeepAlivePool &get_keep_alive_pool()
+{
+    static KeepAlivePool *pool = []() -> KeepAlivePool * {
+        // get dpctl_capi to prevent nullptr return
+        static_cast<void>(dpctl_capi::get());
+
+        return static_cast<KeepAlivePool *>(KeepAlivePool_Get());
+    }();
+
+    if (!pool) {
+        throw std::runtime_error("Could not create dpctl's keep-alive pool");
+    }
+    return *pool;
+}
+
 } // namespace detail
 } // namespace dpctl
 
@@ -828,14 +854,16 @@ sycl::event keep_args_alive(sycl::queue &q,
     }
 
     if (n_usm_owners_held > 0 || n_objects_held > 0) {
-        dpctl::detail::KeepAlivePool::get().submit(
+        dpctl::detail::get_keep_alive_pool().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) {
+                // if the main thread has not finalized the interpreter yet
+                if (n_objects_held > 0 &&
+                    dpctl::detail::interpreter_is_live()) {
                     py::gil_scoped_acquire acquire;
 
                     for (std::size_t i = 0; i < n_objects_held; ++i) {

scripts/_build_helper.py

@@ -93,6 +93,7 @@ def make_cmake_args(
     glog: bool = False,
     verbose: bool = False,
     other_opts: str = None,
+    keep_alive_pool_size: int = None,
 ):
     args = [
         f"-DCMAKE_C_COMPILER:PATH={c_compiler}" if c_compiler else "",
@@ -101,6 +102,10 @@ def make_cmake_args(
         f"-DDPCTL_ENABLE_GLOG:BOOL={'ON' if glog else 'OFF'}",
     ]
 
+    if keep_alive_pool_size is not None:
+        args.append(
+            f"-DDPCTL_KEEP_ALIVE_POOL_SIZE:STRING={keep_alive_pool_size}"
+        )
     if verbose:
         args.append("-DCMAKE_VERBOSE_MAKEFILE:BOOL=ON")
     if other_opts: