Added SyclQueue._submit_keep_args_alive method

Usage:
   q = dpctl.SyclQueue()
   ...
   e = q.submit(krn, args, ranges)
   ht_e = q._submit_keep_args_alive(args, [e])
   ....
   ht_e.wait()

Author: oleksandr-pavlyk

dpctl/_host_task_util.hpp

@@ -2,7 +2,7 @@
 //
 //                      Data Parallel Control (dpctl)
 //
-// Copyright 2020-2022 Intel Corporation
+// Copyright 2020-2023 Intel Corporation
 //
 // Licensed under the Apache License, Version 2.0 (the "License");
 // you may not use this file except in compliance with the License.
@@ -31,28 +31,27 @@
 
 #include "Python.h"
 #include "syclinterface/dpctl_data_types.h"
+#include "syclinterface/dpctl_sycl_type_casters.hpp"
 #include <CL/sycl.hpp>
 
-int async_dec_ref(DPCTLSyclQueueRef QRef,
-                  PyObject **obj_array,
-                  size_t obj_array_size,
-                  DPCTLSyclEventRef *ERefs,
-                  size_t nERefs)
+DPCTLSyclEventRef async_dec_ref(DPCTLSyclQueueRef QRef,
+                                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 = reinterpret_cast<sycl::queue *>(QRef);
+    sycl::queue *q = unwrap<sycl::queue>(QRef);
 
-    std::vector<PyObject *> obj_vec;
-    obj_vec.reserve(obj_array_size);
-    for (size_t obj_id = 0; obj_id < obj_array_size; ++obj_id) {
-        obj_vec.push_back(obj_array[obj_id]);
-    }
+    std::vector<PyObject *> obj_vec(obj_array, obj_array + obj_array_size);
 
     try {
-        q->submit([&](sycl::handler &cgh) {
-            for (size_t ev_id = 0; ev_id < nERefs; ++ev_id) {
-                cgh.depends_on(
-                    *(reinterpret_cast<sycl::event *>(ERefs[ev_id])));
+        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]() {
                 // if the main thread has not finilized the interpreter yet
@@ -66,9 +65,21 @@ int async_dec_ref(DPCTLSyclQueueRef QRef,
                 }
             });
         });
+
+        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) {
-        return 1;
+        constexpr int result_std_exception = 1;
+
+        *status = result_std_exception;
+        return nullptr;
     }
 
-    return 0;
+    constexpr int result_other_abnormal = 2;
+
+    *status = result_other_abnormal;
+    return nullptr;
 }

dpctl/_sycl_queue.pxd

@@ -72,6 +72,11 @@ cdef public api class SyclQueue (_SyclQueue) [
     cpdef SyclContext get_sycl_context(self)
     cpdef SyclDevice get_sycl_device(self)
     cdef  DPCTLSyclQueueRef get_queue_ref(self)
+    cpdef SyclEvent _submit_keep_args_alive(
+        self,
+        object args,
+        list dEvents
+    )
     cpdef SyclEvent submit(
         self,
         SyclKernel kernel,

dpctl/_sycl_queue.pyx

@@ -73,7 +73,7 @@ import logging
 
 
 cdef extern from "_host_task_util.hpp":
-    int async_dec_ref(DPCTLSyclQueueRef, PyObject **, size_t, DPCTLSyclEventRef *, size_t) nogil
+    DPCTLSyclEventRef async_dec_ref(DPCTLSyclQueueRef, PyObject **, size_t, DPCTLSyclEventRef *, size_t, int *) nogil
 
 
 __all__ = [
@@ -716,6 +716,79 @@ cdef class SyclQueue(_SyclQueue):
         """
         return <size_t>self._queue_ref
 
+
+    cpdef SyclEvent _submit_keep_args_alive(
+        self,
+        object args,
+        list dEvents
+    ):
+        """ SyclQueue._submit_keep_args_alive(args, events)
+
+        Keeps objects in `args` alive until tasks associated with events
+        complete.
+
+        Args:
+            args(object): Python object to keep alive.
+               Typically a tuple with arguments to offloaded tasks
+            events(Tuple[dpctl.SyclEvent]): Gating events
+               The list or tuple of events associated with tasks
+               working on Python objects collected in `args`.
+        Returns:
+            dpctl.SyclEvent
+               The event associated with the submission of host task.
+
+        Increments reference count of `args` and schedules asynchronous
+        ``host_task`` to decrement the count once dependent events are
+        complete.
+
+        N.B.: The `host_task` attempts to acquire Python GIL, and it is
+        known to be unsafe during interpreter shudown sequence. It is
+        thus strongly advised to ensure that all submitted `host_task`
+        complete before the end of the Python script.
+        """
+        cdef size_t nDE = len(dEvents)
+        cdef DPCTLSyclEventRef *depEvents = NULL
+        cdef PyObject *args_raw = NULL
+        cdef DPCTLSyclEventRef htERef = NULL
+        cdef int status = -1
+
+        # Create the array of dependent events if any
+        if nDE > 0:
+            depEvents = (
+                <DPCTLSyclEventRef*>malloc(nDE*sizeof(DPCTLSyclEventRef))
+            )
+            if not depEvents:
+                raise MemoryError()
+            else:
+                for idx, de in enumerate(dEvents):
+                    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)
+
+        # schedule decrement
+        args_raw = <PyObject *>args
+
+        htERef = async_dec_ref(
+            self.get_queue_ref(),
+            &args_raw, 1,
+            depEvents, nDE, &status
+        )
+
+        free(depEvents)
+        if (status != 0):
+            with nogil: DPCTLEvent_Wait(htERef)
+            raise RuntimeError("Could not submit keep_args_alive host_task")
+
+        return SyclEvent._create(htERef)
+
+
     cpdef SyclEvent submit(
         self,
         SyclKernel kernel,
@@ -728,13 +801,14 @@ cdef class SyclQueue(_SyclQueue):
         cdef _arg_data_type *kargty = NULL
         cdef DPCTLSyclEventRef *depEvents = NULL
         cdef DPCTLSyclEventRef Eref = NULL
+        cdef DPCTLSyclEventRef htEref = NULL
         cdef int ret = 0
         cdef size_t gRange[3]
         cdef size_t lRange[3]
         cdef size_t nGS = len(gS)
         cdef size_t nLS = len(lS) if lS is not None else 0
         cdef size_t nDE = len(dEvents) if dEvents is not None else 0
-        cdef PyObject **arg_objects = NULL
+        cdef PyObject *args_raw = NULL
         cdef ssize_t i = 0
 
         # Allocate the arrays to be sent to DPCTLQueue_Submit
@@ -758,7 +832,15 @@ cdef class SyclQueue(_SyclQueue):
                 raise MemoryError()
             else:
                 for idx, de in enumerate(dEvents):
-                    depEvents[idx] = (<SyclEvent>de).get_event_ref()
+                    if isinstance(de, SyclEvent):
+                        depEvents[idx] = (<SyclEvent>de).get_event_ref()
+                    else:
+                        free(kargs)
+                        free(kargty)
+                        free(depEvents)
+                        raise TypeError(
+                            "A sequence of dpctl.SyclEvent is expected"
+                        )
 
         # populate the args and argstype arrays
         ret = self._populate_args(args, kargs, kargty)
@@ -836,22 +918,23 @@ cdef class SyclQueue(_SyclQueue):
             raise SyclKernelSubmitError(
                 "Kernel submission to Sycl queue failed."
             )
-        # increment reference counts to each argument
-        arg_objects = <PyObject **>malloc(len(args) * sizeof(PyObject *))
-        for i in range(len(args)):
-            arg_objects[i] = <PyObject *>(args[i])
-            Py_INCREF(<object> arg_objects[i])
+        # increment reference counts to list of arguments
+        Py_INCREF(args)
 
         # schedule decrement
-        if async_dec_ref(self.get_queue_ref(), arg_objects, len(args), &Eref, 1):
+        args_raw = <PyObject *>args
+
+        ret = -1
+        htERef = async_dec_ref(self.get_queue_ref(), &args_raw, 1, &Eref, 1, &ret)
+        if ret:
             # async task submission failed, decrement ref counts and wait
-            for i in range(len(args)):
-                arg_objects[i] = <PyObject *>(args[i])
-                Py_DECREF(<object> arg_objects[i])
-            with nogil: DPCTLEvent_Wait(Eref)
+            Py_DECREF(args)
+            with nogil:
+                 DPCTLEvent_Wait(Eref)
+                 DPCTLEvent_Wait(htERef)
 
-        # free memory
-        free(arg_objects)
+        # we are not returning host-task event at the moment
+        DPCTLEvent_Delete(htERef)
 
         return SyclEvent._create(Eref)