Accumulate implementations API acquired host_task_events argument

This is `std::vector<sycl::event>` passed by reference to collect
events associated with host_task submissions.

The synchronizing call `mask_positions` is releasing GIL before
wait on these events is called.

It is either this, or accumulated host tasks must be returned to
the user.

Author: oleksandr-pavlyk

dpctl/tensor/libtensor/include/kernels/accumulators.hpp

@@ -112,7 +112,8 @@ sycl::event inclusive_scan_rec(sycl::queue &exec_q,
                                size_t s1,
                                IndexerT indexer,
                                TransformerT transformer,
-                               std::vector<sycl::event> const &depends = {})
+                               std::vector<sycl::event> &host_tasks,
+                               const std::vector<sycl::event> &depends = {})
 {
     size_t n_groups = ceiling_quotient(n_elems, n_wi * wg_size);
 
@@ -188,7 +189,7 @@ sycl::event inclusive_scan_rec(sycl::queue &exec_q,
         auto e2 = inclusive_scan_rec<outputT, outputT, n_wi, NoOpIndexer,
                                      decltype(_no_op_transformer)>(
             exec_q, n_groups - 1, wg_size, output, temp, chunk_size - 1,
-            chunk_size, _no_op_indexer, _no_op_transformer,
+            chunk_size, _no_op_indexer, _no_op_transformer, host_tasks,
             {inc_scan_phase1_ev});
 
         // output[ chunk_size * (i + 1) + j] += temp[i]
@@ -209,8 +210,9 @@ sycl::event inclusive_scan_rec(sycl::queue &exec_q,
             const auto &ctx = exec_q.get_context();
             cgh.host_task([ctx, temp]() { sycl::free(temp, ctx); });
         });
+        host_tasks.push_back(e4);
 
-        out_event = std::move(e4);
+        out_event = std::move(e3);
     }
 
     return out_event;
@@ -221,14 +223,16 @@ typedef size_t (*accumulate_contig_impl_fn_ptr_t)(
     size_t,
     const char *,
     char *,
-    std::vector<sycl::event> const &);
+    std::vector<sycl::event> &,
+    const std::vector<sycl::event> &);
 
 template <typename maskT, typename cumsumT, typename transformerT>
 size_t accumulate_contig_impl(sycl::queue &q,
                               size_t n_elems,
                               const char *mask,
                               char *cumsum,
-                              std::vector<sycl::event> const &depends = {})
+                              std::vector<sycl::event> &host_tasks,
+                              const std::vector<sycl::event> &depends = {})
 {
     constexpr int n_wi = 8;
     const maskT *mask_data_ptr = reinterpret_cast<const maskT *>(mask);
@@ -242,7 +246,7 @@ size_t accumulate_contig_impl(sycl::queue &q,
         inclusive_scan_rec<maskT, cumsumT, n_wi, decltype(flat_indexer),
                            decltype(non_zero_indicator)>(
             q, n_elems, wg_size, mask_data_ptr, cumsum_data_ptr, 0, 1,
-            flat_indexer, non_zero_indicator, depends);
+            flat_indexer, non_zero_indicator, host_tasks, depends);
 
     cumsumT *last_elem = cumsum_data_ptr + (n_elems - 1);
 
@@ -307,7 +311,8 @@ typedef size_t (*accumulate_strided_impl_fn_ptr_t)(
     int,
     const py::ssize_t *,
     char *,
-    std::vector<sycl::event> const &);
+    std::vector<sycl::event> &,
+    const std::vector<sycl::event> &);
 
 template <typename maskT, typename cumsumT, typename transformerT>
 size_t accumulate_strided_impl(sycl::queue &q,
@@ -316,7 +321,8 @@ size_t accumulate_strided_impl(sycl::queue &q,
                                int nd,
                                const py::ssize_t *shape_strides,
                                char *cumsum,
-                               std::vector<sycl::event> const &depends = {})
+                               std::vector<sycl::event> &host_tasks,
+                               const std::vector<sycl::event> &depends = {})
 {
     constexpr int n_wi = 8;
     const maskT *mask_data_ptr = reinterpret_cast<const maskT *>(mask);
@@ -330,7 +336,7 @@ size_t accumulate_strided_impl(sycl::queue &q,
         inclusive_scan_rec<maskT, cumsumT, n_wi, decltype(strided_indexer),
                            decltype(non_zero_indicator)>(
             q, n_elems, wg_size, mask_data_ptr, cumsum_data_ptr, 0, 1,
-            strided_indexer, non_zero_indicator, depends);
+            strided_indexer, non_zero_indicator, host_tasks, depends);
 
     cumsumT *last_elem = cumsum_data_ptr + (n_elems - 1);
 

dpctl/tensor/libtensor/source/accumulators.cpp

@@ -150,12 +150,20 @@ size_t py_mask_positions(const dpctl::tensor::usm_ndarray &mask,
 
     const bool use_i32 = (cumsum_typeid == int32_typeid);
 
+    std::vector<sycl::event> host_task_events;
+
     if (mask.is_c_contiguous()) {
         auto fn = (use_i32)
                       ? mask_positions_contig_i32_dispatch_vector[mask_typeid]
                       : mask_positions_contig_i64_dispatch_vector[mask_typeid];
 
-        return fn(exec_q, mask_size, mask_data, cumsum_data, depends);
+        size_t total_set = fn(exec_q, mask_size, mask_data, cumsum_data,
+                              host_task_events, depends);
+        {
+            py::gil_scoped_release release;
+            sycl::event::wait(host_task_events);
+        }
+        return total_set;
     }
 
     const py::ssize_t *shape = mask.get_shape_raw();
@@ -175,7 +183,6 @@ size_t py_mask_positions(const dpctl::tensor::usm_ndarray &mask,
     auto strided_fn =
         (use_i32) ? mask_positions_strided_i32_dispatch_vector[mask_typeid]
                   : mask_positions_strided_i64_dispatch_vector[mask_typeid];
-    std::vector<sycl::event> host_task_events;
 
     using dpctl::tensor::offset_utils::device_allocate_and_pack;
     const auto &ptr_size_event_tuple = device_allocate_and_pack<py::ssize_t>(
@@ -189,7 +196,10 @@ size_t py_mask_positions(const dpctl::tensor::usm_ndarray &mask,
 
     if (2 * static_cast<size_t>(nd) != std::get<1>(ptr_size_event_tuple)) {
         copy_shape_ev.wait();
-        sycl::event::wait(host_task_events);
+        {
+            py::gil_scoped_release release;
+            sycl::event::wait(host_task_events);
+        }
         sycl::free(shape_strides, exec_q);
         throw std::runtime_error("Unexpected error");
     }
@@ -200,10 +210,14 @@ size_t py_mask_positions(const dpctl::tensor::usm_ndarray &mask,
     dependent_events.insert(dependent_events.end(), depends.begin(),
                             depends.end());
 
-    size_t total_set = strided_fn(exec_q, mask_size, mask_data, nd,
-                                  shape_strides, cumsum_data, dependent_events);
+    size_t total_set =
+        strided_fn(exec_q, mask_size, mask_data, nd, shape_strides, cumsum_data,
+                   host_task_events, dependent_events);
 
-    sycl::event::wait(host_task_events);
+    {
+        py::gil_scoped_release release;
+        sycl::event::wait(host_task_events);
+    }
     sycl::free(shape_strides, exec_q);
 
     return total_set;
@@ -283,14 +297,22 @@ size_t py_cumsum_1d(const dpctl::tensor::usm_ndarray &src,
             "Cumulative sum array must have int64 data-type.");
     }
 
+    std::vector<sycl::event> host_task_events;
+
     if (src.is_c_contiguous()) {
         auto fn = cumsum_1d_contig_dispatch_vector[src_typeid];
         if (fn == nullptr) {
             throw std::runtime_error(
                 "this cumsum requires integer type, got src_typeid=" +
                 std::to_string(src_typeid));
         }
-        return fn(exec_q, src_size, src_data, cumsum_data, depends);
+        size_t total = fn(exec_q, src_size, src_data, cumsum_data,
+                          host_task_events, depends);
+        {
+            py::gil_scoped_release release;
+            sycl::event::wait(host_task_events);
+        }
+        return total;
     }
 
     const py::ssize_t *shape = src.get_shape_raw();
@@ -313,7 +335,6 @@ size_t py_cumsum_1d(const dpctl::tensor::usm_ndarray &src,
             "this cumsum requires integer type, got src_typeid=" +
             std::to_string(src_typeid));
     }
-    std::vector<sycl::event> host_task_events;
 
     using dpctl::tensor::offset_utils::device_allocate_and_pack;
     const auto &ptr_size_event_tuple = device_allocate_and_pack<py::ssize_t>(
@@ -339,9 +360,12 @@ size_t py_cumsum_1d(const dpctl::tensor::usm_ndarray &src,
                             depends.end());
 
     size_t total = strided_fn(exec_q, src_size, src_data, nd, shape_strides,
-                              cumsum_data, dependent_events);
+                              cumsum_data, host_task_events, dependent_events);
 
-    sycl::event::wait(host_task_events);
+    {
+        py::gil_scoped_release release;
+        sycl::event::wait(host_task_events);
+    }
     sycl::free(shape_strides, exec_q);
 
     return total;