gemv function returns a pair of events, first for host task, second for MKL computation

Author: oleksandr-pavlyk

examples/pybind11/onemkl_gemv/sycl_gemm/_onemkl.cpp

@@ -16,11 +16,41 @@ namespace py = pybind11;
      UsmNDArray_GetQueueRef
  */
 
-sycl::event gemv(sycl::queue q,
-                 py::object matrix,
-                 py::object vector,
-                 py::object result,
-                 const std::vector<sycl::event> &depends = {})
+sycl::event keep_args_alive(sycl::queue q,
+                            py::object o1,
+                            py::object o2,
+                            py::object o3,
+                            const std::vector<sycl::event> &depends = {})
+{
+    sycl::event ht_event = q.submit([&](sycl::handler &cgh) {
+        cgh.depends_on(depends);
+        std::shared_ptr<py::handle> shp1 = std::make_shared<py::handle>(o1);
+        std::shared_ptr<py::handle> shp2 = std::make_shared<py::handle>(o2);
+        std::shared_ptr<py::handle> shp3 = std::make_shared<py::handle>(o3);
+        shp1->inc_ref();
+        shp2->inc_ref();
+        shp3->inc_ref();
+        cgh.host_task([=]() {
+            bool guard = (Py_IsInitialized() && !_Py_IsFinalizing());
+            if (guard) {
+                PyGILState_STATE gstate;
+                gstate = PyGILState_Ensure();
+                shp1->dec_ref();
+                shp2->dec_ref();
+                shp3->dec_ref();
+                PyGILState_Release(gstate);
+            }
+        });
+    });
+    return ht_event;
+}
+
+std::pair<sycl::event, sycl::event>
+gemv(sycl::queue q,
+     py::object matrix,
+     py::object vector,
+     py::object result,
+     const std::vector<sycl::event> &depends = {})
 {
     PyObject *m_src = matrix.ptr();
     if (!PyObject_TypeCheck(m_src, &PyUSMArrayType)) {
@@ -131,7 +161,9 @@ sycl::event gemv(sycl::queue q,
         throw std::runtime_error("Type dispatch ran into trouble.");
     }
 
-    return res_ev;
+    sycl::event ht_event = keep_args_alive(q, matrix, vector, result, {res_ev});
+
+    return std::make_pair(ht_event, res_ev);
 }
 
 PYBIND11_MODULE(_onemkl, m)

examples/pybind11/onemkl_gemv/tests/test_gemm.py

@@ -15,7 +15,7 @@ def test_gemv():
     r = dpt.empty((5,), dtype="d", sycl_queue=q)
     M = dpt.asarray(Mnp, sycl_queue=q)
     v = dpt.asarray(vnp, sycl_queue=q)
-    ev = gemv(M.sycl_queue, M, v, r, [])
-    ev.wait()
+    hev, ev = gemv(M.sycl_queue, M, v, r, [])
+    hev.wait()
     rnp = dpt.asnumpy(r)
     assert np.allclose(rnp, Mnp @ vnp)