[cublas] introduce onemkl_cublas_host_task

[cublas] move dpc++ internal headers into cublas_task.hpp

Author: sbalint98

src/blas/backends/cublas/cublas_batch.cpp

@@ -16,9 +16,9 @@
 *  limitations under the License.
 *
 **************************************************************************/
-#include <CL/sycl/detail/pi.hpp>
 #include "cublas_helper.hpp"
-#include "cublas_scope_handle.hpp"
+#include "cublas_task.hpp"
+
 #include "oneapi/mkl/exceptions.hpp"
 #include "oneapi/mkl/blas/detail/cublas/onemkl_blas_cublas.hpp"
 
@@ -42,12 +42,12 @@ inline void gemm_batch(Func func, cl::sycl::queue &queue, transpose transa, tran
         auto a_acc = a.template get_access<cl::sycl::access::mode::read>(cgh);
         auto b_acc = b.template get_access<cl::sycl::access::mode::read>(cgh);
         auto c_acc = c.template get_access<cl::sycl::access::mode::read_write>(cgh);
-        cgh.interop_task([=](cl::sycl::interop_handler ih) {
-            auto sc = CublasScopedContextHandler(queue);
+        onemkl_cublas_host_task(cgh, queue,[=](CublasScopedContextHandler sc) {
             auto handle = sc.get_handle(queue);
-            auto a_ = sc.get_mem<cuDataType *>(ih, a_acc);
-            auto b_ = sc.get_mem<cuDataType *>(ih, b_acc);
-            auto c_ = sc.get_mem<cuDataType *>(ih, c_acc);
+
+            auto a_ = sc.get_mem<cuDataType *>(a_acc);
+            auto b_ = sc.get_mem<cuDataType *>(b_acc);
+            auto c_ = sc.get_mem<cuDataType *>(c_acc);
             cublasStatus_t err;
             CUBLAS_ERROR_FUNC(func, err, handle, get_cublas_operation(transa),
                               get_cublas_operation(transb), m, n, k, (cuDataType *)&alpha, a_, lda,
@@ -122,9 +122,9 @@ inline cl::sycl::event gemm_batch(Func func, cl::sycl::queue &queue, transpose t
         for (int64_t i = 0; i < num_events; i++) {
             cgh.depends_on(dependencies[i]);
         }
-        cgh.interop_task([=](cl::sycl::interop_handler ih) {
-            auto sc = CublasScopedContextHandler(queue);
+        onemkl_cublas_host_task(cgh, queue,[=](CublasScopedContextHandler sc) {
             auto handle = sc.get_handle(queue);
+
             auto a_ = reinterpret_cast<const cuDataType *>(a);
             auto b_ = reinterpret_cast<const cuDataType *>(b);
             auto c_ = reinterpret_cast<cuDataType *>(c);
@@ -170,9 +170,9 @@ inline cl::sycl::event gemm_batch(Func func, cl::sycl::queue &queue, transpose *
         for (int64_t i = 0; i < num_events; i++) {
             cgh.depends_on(dependencies[i]);
         }
-        cgh.interop_task([=](cl::sycl::interop_handler ih) {
-            auto sc = CublasScopedContextHandler(queue);
+        onemkl_cublas_host_task(cgh, queue,[=](CublasScopedContextHandler sc) {
             auto handle = sc.get_handle(queue);
+
             int64_t offset = 0;
             cublasStatus_t err;
             for (int64_t i = 0; i < group_count; i++) {

src/blas/backends/cublas/cublas_extensions.cpp

@@ -16,9 +16,9 @@
 *  limitations under the License.
 *
 **************************************************************************/
-#include <CL/sycl/detail/pi.hpp>
 #include "cublas_helper.hpp"
-#include "cublas_scope_handle.hpp"
+#include "cublas_task.hpp"
+
 #include "oneapi/mkl/exceptions.hpp"
 #include "oneapi/mkl/blas/detail/cublas/onemkl_blas_cublas.hpp"
 

src/blas/backends/cublas/cublas_level1.cpp

@@ -48,7 +48,7 @@ cublas_handle::~cublas_handle() noexcept(false) {
  */
 thread_local cublas_handle CublasScopedContextHandler::handle_helper = cublas_handle{};
 
-CublasScopedContextHandler::CublasScopedContextHandler(cl::sycl::queue queue) {
+CublasScopedContextHandler::CublasScopedContextHandler(cl::sycl::queue queue, cl::sycl::interop_handler& ih): ih(ih){
     placedContext_ = queue.get_context();
     auto device = queue.get_device();
     auto desired = cl::sycl::get_native<cl::sycl::backend::cuda>(placedContext_);

src/blas/backends/cublas/cublas_level2.cpp

@@ -68,12 +68,13 @@ class CublasScopedContextHandler {
     CUcontext original_;
     cl::sycl::context placedContext_;
     bool needToRecover_;
+    cl::sycl::interop_handler& ih;
     static thread_local cublas_handle handle_helper;
     CUstream get_stream(const cl::sycl::queue &queue);
     cl::sycl::context get_context(const cl::sycl::queue &queue);
 
 public:
-    CublasScopedContextHandler(cl::sycl::queue queue);
+    CublasScopedContextHandler(cl::sycl::queue queue, cl::sycl::interop_handler& ih);
 
     ~CublasScopedContextHandler() noexcept(false);
     /**
@@ -87,7 +88,7 @@ class CublasScopedContextHandler {
     // This is a work-around function for reinterpret_casting the memory. This
     // will be fixed when SYCL-2020 has been implemented for Pi backend.
     template <typename T, typename U>
-    inline T get_mem(cl::sycl::interop_handler ih, U acc) {
+    inline T get_mem(U acc) {
         CUdeviceptr cudaPtr = ih.get_mem<cl::sycl::backend::cuda>(acc);
         return reinterpret_cast<T>(cudaPtr);
     }

src/blas/backends/cublas/cublas_level3.cpp

@@ -0,0 +1,33 @@
+#ifndef _MKL_BLAS_CUBLAS_TASK_HPP_
+#define _MKL_BLAS_CUBLAS_TASK_HPP_
+#include <cublas_v2.h>
+#include <cuda.h>
+#include <complex>
+#include <CL/sycl.hpp>
+#include "oneapi/mkl/types.hpp"
+#include "cublas_scope_handle.hpp"
+#include <CL/sycl/detail/pi.hpp>
+
+namespace oneapi {
+namespace mkl {
+namespace blas {
+namespace cublas {
+
+template <typename H, typename F>
+static inline auto host_task_internal(H &cgh, cl::sycl::queue queue, F f) -> decltype(cgh.interop_task(f)) {
+    cgh.interop_task([f, queue](cl::sycl::interop_handler ih){
+        auto sc = CublasScopedContextHandler(queue, ih);
+        f(sc);
+    });
+}
+
+template <typename H, typename F>
+static inline void onemkl_cublas_host_task(H &cgh, cl::sycl::queue queue, F f) {
+    (void)host_task_internal(cgh, queue, f);
+}
+
+} // namespace cublas
+} // namespace blas
+} // namespace mkl
+} // namespace oneapi
+#endif // _MKL_BLAS_CUBLAS_TASK_HPP_