Pass queue from cython to backend for matmul (#1063)

densmirn · Alexander-Makaryev · web-flow · commit 180a69800b7e · 2022-01-28T15:56:21.000+03:00
* Pass queue from cython to backend for matmul

* Add DPNP_FN_MATMUL_EXT to be compatible with DPPY

* minor fix in DPNPFuncName enum

* Move parameter ueue to the first place in dpnp_matmul_c

* Change void* to DPCTLSyclQueueRef + add test to check queue of result

* add test to check queue of result

* Limit list of SYCL devices for testing

Co-authored-by: Alexander-Makaryev &lt;alexander.makaryev@gmail.com&gt;
diff --git a/dpnp/backend/include/dpnp_iface.hpp b/dpnp/backend/include/dpnp_iface.hpp
@@ -58,6 +58,8 @@ typedef ssize_t shape_elem_type;
 #include "dpnp_iface_fft.hpp"
 #include "dpnp_iface_random.hpp"
 
+#include <dpctl_sycl_interface.h>
+
 /**
  * @defgroup BACKEND_API Backend C++ library interface API
  * @{
@@ -196,6 +198,47 @@ INP_DLLEXPORT void dpnp_full_c(void* array_in, void* result, const size_t size);
 template <typename _DataType>
 INP_DLLEXPORT void dpnp_full_like_c(void* array_in, void* result, size_t size);
 
+/**
+ * @ingroup BACKEND_API
+ * @brief Matrix multiplication.
+ *
+ * Matrix multiplication procedure.
+ *
+ * @param [in]  q_ref               Reference to SYCL queue.
+ * @param [out] result_out          Output array.
+ * @param [in]  result_size         Size of output array.
+ * @param [in]  result_ndim         Number of output array dimensions.
+ * @param [in]  result_shape        Shape of output array.
+ * @param [in]  result_strides      Strides of output array.
+ * @param [in]  input1_in           First input array.
+ * @param [in]  input1_size         Size of first input array.
+ * @param [in]  input1_ndim         Number of first input array dimensions.
+ * @param [in]  input1_shape        Shape of first input array.
+ * @param [in]  input1_strides      Strides of first input array.
+ * @param [in]  input2_in           Second input array.
+ * @param [in]  input2_size         Size of second input array.
+ * @param [in]  input2_ndim         Number of second input array dimensions.
+ * @param [in]  input2_shape        Shape of second input array.
+ * @param [in]  input2_strides      Strides of second input array.
+ */
+template <typename _DataType>
+INP_DLLEXPORT void dpnp_matmul_c(DPCTLSyclQueueRef q_ref,
+                                 void* result_out,
+                                 const size_t result_size,
+                                 const size_t result_ndim,
+                                 const shape_elem_type* result_shape,
+                                 const shape_elem_type* result_strides,
+                                 const void* input1_in,
+                                 const size_t input1_size,
+                                 const size_t input1_ndim,
+                                 const shape_elem_type* input1_shape,
+                                 const shape_elem_type* input1_strides,
+                                 const void* input2_in,
+                                 const size_t input2_size,
+                                 const size_t input2_ndim,
+                                 const shape_elem_type* input2_shape,
+                                 const shape_elem_type* input2_strides);
+
 /**
  * @ingroup BACKEND_API
  * @brief Matrix multiplication.
diff --git a/dpnp/backend/include/dpnp_iface_fptr.hpp b/dpnp/backend/include/dpnp_iface_fptr.hpp
@@ -131,6 +131,7 @@ enum class DPNPFuncName : size_t
     DPNP_FN_LOG2,                     /**< Used in numpy.log2() implementation  */
     DPNP_FN_LOG1P,                    /**< Used in numpy.log1p() implementation  */
     DPNP_FN_MATMUL,                   /**< Used in numpy.matmul() implementation  */
+    DPNP_FN_MATMUL_EXT,               /**< Used in numpy.matmul() implementation, requires extra parameters  */
     DPNP_FN_MATRIX_RANK,              /**< Used in numpy.linalg.matrix_rank() implementation  */
     DPNP_FN_MAX,                      /**< Used in numpy.max() implementation  */
     DPNP_FN_MAXIMUM,                  /**< Used in numpy.maximum() implementation  */
diff --git a/dpnp/backend/kernels/dpnp_krnl_common.cpp b/dpnp/backend/kernels/dpnp_krnl_common.cpp
@@ -533,7 +533,8 @@ template <typename _KernelNameSpecialization>
 class dpnp_matmul_c_kernel;
 
 template <typename _DataType>
-void dpnp_matmul_c(void* result_out,
+void dpnp_matmul_c(DPCTLSyclQueueRef q_ref,
+                   void* result_out,
                    const size_t result_size,
                    const size_t result_ndim,
                    const shape_elem_type* result_shape,
@@ -569,13 +570,12 @@ void dpnp_matmul_c(void* result_out,
         return;
     }
 
+    sycl::queue q = *(reinterpret_cast<sycl::queue*>(q_ref));
     sycl::event event;
-    DPNPC_ptr_adapter<_DataType> input1_ptr(input1_in, size_m * size_k);
-    DPNPC_ptr_adapter<_DataType> input2_ptr(input2_in, size_k * size_n);
-    DPNPC_ptr_adapter<_DataType> result_ptr(result_out, size_m * size_n, false, true);
-    _DataType* array_1 = input1_ptr.get_ptr();
-    _DataType* array_2 = input2_ptr.get_ptr();
-    _DataType* result = result_ptr.get_ptr();
+
+    _DataType* array_1 = reinterpret_cast<_DataType*>(const_cast<void*>(input1_in));
+    _DataType* array_2 = reinterpret_cast<_DataType*>(const_cast<void*>(input2_in));
+    _DataType* result = reinterpret_cast<_DataType*>(result_out);
 
     if constexpr (std::is_same<_DataType, double>::value || std::is_same<_DataType, float>::value)
     {
@@ -584,7 +584,7 @@ void dpnp_matmul_c(void* result_out,
         const std::int64_t ldb = std::max<size_t>(1UL, size_n); // First dimensions of array_2
         const std::int64_t ldc = std::max<size_t>(1UL, size_n); // Fast dimensions of result
 
-        event = mkl_blas::gemm(DPNP_QUEUE,
+        event = mkl_blas::gemm(q,
                                oneapi::mkl::transpose::nontrans,
                                oneapi::mkl::transpose::nontrans,
                                size_n,
@@ -632,11 +632,70 @@ void dpnp_matmul_c(void* result_out,
             cgh.parallel_for<class dpnp_matmul_c_kernel<_DataType>>(gws, kernel_parallel_for_func);
         };
 
-        event = DPNP_QUEUE.submit(kernel_func);
+        event = q.submit(kernel_func);
     }
     event.wait();
 }
 
+template <typename _DataType>
+void dpnp_matmul_c(void* result_out,
+                   const size_t result_size,
+                   const size_t result_ndim,
+                   const shape_elem_type* result_shape,
+                   const shape_elem_type* result_strides,
+                   const void* input1_in,
+                   const size_t input1_size,
+                   const size_t input1_ndim,
+                   const shape_elem_type* input1_shape,
+                   const shape_elem_type* input1_strides,
+                   const void* input2_in,
+                   const size_t input2_size,
+                   const size_t input2_ndim,
+                   const shape_elem_type* input2_shape,
+                   const shape_elem_type* input2_strides)
+{
+    DPCTLSyclQueueRef q_ref = reinterpret_cast<DPCTLSyclQueueRef>(&DPNP_QUEUE);
+    dpnp_matmul_c<_DataType>(q_ref,
+                             result_out, result_size, result_ndim, result_shape, result_strides,
+                             input1_in, input1_size, input1_ndim, input1_shape, input1_strides,
+                             input2_in, input2_size, input2_ndim, input2_shape, input2_strides);
+}
+
+template <typename _DataType>
+void (*dpnp_matmul_default_c)(void*,
+                              const size_t,
+                              const size_t,
+                              const shape_elem_type*,
+                              const shape_elem_type*,
+                              const void*,
+                              const size_t,
+                              const size_t,
+                              const shape_elem_type*,
+                              const shape_elem_type*,
+                              const void*,
+                              const size_t,
+                              const size_t,
+                              const shape_elem_type*,
+                              const shape_elem_type*) = dpnp_matmul_c<_DataType>;
+
+template <typename _DataType>
+void (*dpnp_matmul_ext_c)(DPCTLSyclQueueRef,
+                          void*,
+                          const size_t,
+                          const size_t,
+                          const shape_elem_type*,
+                          const shape_elem_type*,
+                          const void*,
+                          const size_t,
+                          const size_t,
+                          const shape_elem_type*,
+                          const shape_elem_type*,
+                          const void*,
+                          const size_t,
+                          const size_t,
+                          const shape_elem_type*,
+                          const shape_elem_type*) = dpnp_matmul_c<_DataType>;
+
 void func_map_init_linalg(func_map_t& fmap)
 {
     fmap[DPNPFuncName::DPNP_FN_ASTYPE][eft_BLN][eft_BLN] = {eft_BLN, (void*)dpnp_astype_c<bool, bool>};
@@ -702,10 +761,15 @@ void func_map_init_linalg(func_map_t& fmap)
     fmap[DPNPFuncName::DPNP_FN_INITVAL][eft_DBL][eft_DBL] = {eft_DBL, (void*)dpnp_initval_c<double>};
     fmap[DPNPFuncName::DPNP_FN_INITVAL][eft_C128][eft_C128] = {eft_C128, (void*)dpnp_initval_c<std::complex<double>>};
 
-    fmap[DPNPFuncName::DPNP_FN_MATMUL][eft_INT][eft_INT] = {eft_INT, (void*)dpnp_matmul_c<int32_t>};
-    fmap[DPNPFuncName::DPNP_FN_MATMUL][eft_LNG][eft_LNG] = {eft_LNG, (void*)dpnp_matmul_c<int64_t>};
-    fmap[DPNPFuncName::DPNP_FN_MATMUL][eft_FLT][eft_FLT] = {eft_FLT, (void*)dpnp_matmul_c<float>};
-    fmap[DPNPFuncName::DPNP_FN_MATMUL][eft_DBL][eft_DBL] = {eft_DBL, (void*)dpnp_matmul_c<double>};
+    fmap[DPNPFuncName::DPNP_FN_MATMUL][eft_INT][eft_INT] = {eft_INT, (void*)dpnp_matmul_default_c<int32_t>};
+    fmap[DPNPFuncName::DPNP_FN_MATMUL][eft_LNG][eft_LNG] = {eft_LNG, (void*)dpnp_matmul_default_c<int64_t>};
+    fmap[DPNPFuncName::DPNP_FN_MATMUL][eft_FLT][eft_FLT] = {eft_FLT, (void*)dpnp_matmul_default_c<float>};
+    fmap[DPNPFuncName::DPNP_FN_MATMUL][eft_DBL][eft_DBL] = {eft_DBL, (void*)dpnp_matmul_default_c<double>};
+
+    fmap[DPNPFuncName::DPNP_FN_MATMUL_EXT][eft_INT][eft_INT] = {eft_INT, (void*)dpnp_matmul_ext_c<int32_t>};
+    fmap[DPNPFuncName::DPNP_FN_MATMUL_EXT][eft_LNG][eft_LNG] = {eft_LNG, (void*)dpnp_matmul_ext_c<int64_t>};
+    fmap[DPNPFuncName::DPNP_FN_MATMUL_EXT][eft_FLT][eft_FLT] = {eft_FLT, (void*)dpnp_matmul_ext_c<float>};
+    fmap[DPNPFuncName::DPNP_FN_MATMUL_EXT][eft_DBL][eft_DBL] = {eft_DBL, (void*)dpnp_matmul_ext_c<double>};
 
     return;
 }
diff --git a/dpnp/dpnp_algo/dpnp_algo.pxd b/dpnp/dpnp_algo/dpnp_algo.pxd
@@ -106,6 +106,7 @@ cdef extern from "dpnp_iface_fptr.hpp" namespace "DPNPFuncName":  # need this na
         DPNP_FN_LOG1P
         DPNP_FN_LOG2
         DPNP_FN_MATMUL
+        DPNP_FN_MATMUL_EXT
         DPNP_FN_MATRIX_RANK
         DPNP_FN_MAX
         DPNP_FN_MAXIMUM
diff --git a/dpnp/dpnp_algo/dpnp_algo.pyx b/dpnp/dpnp_algo/dpnp_algo.pyx
@@ -39,12 +39,14 @@ import dpnp.config as config
 import dpnp.dpnp_utils as utils_py
 from dpnp.dpnp_array import dpnp_array
 
-import numpy
+cimport dpctl as c_dpctl
 import dpctl
 
 cimport cpython
 cimport dpnp.dpnp_utils as utils
+
 cimport numpy
+import numpy
 
 
 __all__ = [
diff --git a/dpnp/dpnp_algo/dpnp_algo_linearalgebra.pyx b/dpnp/dpnp_algo/dpnp_algo_linearalgebra.pyx
@@ -52,6 +52,13 @@ ctypedef void(*fptr_2in_1out_dot_t)(void * , const size_t, const size_t,
                                     const shape_elem_type *, const shape_elem_type * ,
                                     void * , const size_t, const size_t,
                                     const shape_elem_type *, const shape_elem_type * )
+ctypedef void(*fptr_2in_1out_matmul_t)(c_dpctl.DPCTLSyclQueueRef,
+                                       void * , const size_t, const size_t,
+                                       const shape_elem_type *, const shape_elem_type * ,
+                                       void * , const size_t, const size_t,
+                                       const shape_elem_type *, const shape_elem_type * ,
+                                       void * , const size_t, const size_t,
+                                       const shape_elem_type *, const shape_elem_type * )
 
 cpdef utils.dpnp_descriptor dpnp_dot(utils.dpnp_descriptor in_array1, utils.dpnp_descriptor in_array2):
 
@@ -271,7 +278,7 @@ cpdef utils.dpnp_descriptor dpnp_matmul(utils.dpnp_descriptor in_array1, utils.d
     cdef DPNPFuncType param2_type = dpnp_dtype_to_DPNPFuncType(in_array2.dtype)
 
     # get the FPTR data structure
-    cdef DPNPFuncData kernel_data = get_dpnp_function_ptr(DPNP_FN_MATMUL, param1_type, param2_type)
+    cdef DPNPFuncData kernel_data = get_dpnp_function_ptr(DPNP_FN_MATMUL_EXT, param1_type, param2_type)
 
     # ceate result array with type given by FPTR data
     result_sycl_device, result_usm_type, result_sycl_queue = utils.get_common_usm_allocation(in_array1, in_array2)
@@ -284,9 +291,13 @@ cpdef utils.dpnp_descriptor dpnp_matmul(utils.dpnp_descriptor in_array1, utils.d
     if result.size == 0:
         return result
 
-    cdef fptr_2in_1out_dot_t func = <fptr_2in_1out_dot_t > kernel_data.ptr
+    cdef c_dpctl.SyclQueue q = <c_dpctl.SyclQueue> result_sycl_queue
+    cdef c_dpctl.DPCTLSyclQueueRef q_ref = q.get_queue_ref()
+
+    cdef fptr_2in_1out_matmul_t func = <fptr_2in_1out_matmul_t > kernel_data.ptr
     # call FPTR function
-    func(result.get_data(),
+    func(q_ref,
+         result.get_data(),
          result.size,
          result.ndim,
          NULL,  # result_shape
diff --git a/dpnp/dpnp_iface.py b/dpnp/dpnp_iface.py
@@ -186,7 +186,7 @@ def convert_single_elem_array_to_scalar(obj, keepdims=False):
     return obj
 
 
-def get_dpnp_descriptor(ext_obj, copy_when_strides=True):
+def get_dpnp_descriptor(ext_obj, copy_when_strides=True, copy_when_nondefault_queue=True):
     """
     Return True:
       never
@@ -221,6 +221,18 @@ def get_dpnp_descriptor(ext_obj, copy_when_strides=True):
         if ext_obj.strides != shape_offsets or ext_obj_offset != 0:
             ext_obj = array(ext_obj)
 
+    # while dpnp functions are based on DPNP_QUEUE
+    # we need to create a copy on device associated with DPNP_QUEUE
+    # if function get implementation for different queue
+    # then this behavior can be disabled with setting "copy_when_nondefault_queue"
+    arr_obj = unwrap_array(ext_obj)
+    queue = getattr(arr_obj, "sycl_queue", None)
+    if queue is not None and copy_when_nondefault_queue:
+        default_queue = dpctl.SyclQueue()
+        queue_is_default = dpctl.utils.get_execution_queue([queue, default_queue]) is not None
+        if not queue_is_default:
+            ext_obj = array(arr_obj, sycl_queue=default_queue)
+
     dpnp_desc = dpnp_descriptor(ext_obj)
     if dpnp_desc.is_valid:
         return dpnp_desc
diff --git a/dpnp/dpnp_iface_linearalgebra.py b/dpnp/dpnp_iface_linearalgebra.py
@@ -241,8 +241,8 @@ def matmul(x1, x2, out=None, **kwargs):
 
     """
 
-    x1_desc = dpnp.get_dpnp_descriptor(x1)
-    x2_desc = dpnp.get_dpnp_descriptor(x2)
+    x1_desc = dpnp.get_dpnp_descriptor(x1, copy_when_nondefault_queue=False)
+    x2_desc = dpnp.get_dpnp_descriptor(x2, copy_when_nondefault_queue=False)
     if x1_desc and x2_desc and not kwargs:
         if x1_desc.ndim != 2 or x2_desc.ndim != 2:
             pass
diff --git a/dpnp/dpnp_utils/dpnp_algo_utils.pyx b/dpnp/dpnp_utils/dpnp_algo_utils.pyx
@@ -63,6 +63,7 @@ __all__ = [
     "_get_linear_index",
     "normalize_axis",
     "_object_to_tuple",
+    "unwrap_array",
     "use_origin_backend"
 ]
 
diff --git a/setup.py b/setup.py
@@ -40,6 +40,7 @@
 import importlib.machinery as imm  # Python 3 is required
 import sys
 import os
+import dpctl
 import numpy
 
 from setuptools import setup, Extension
@@ -132,7 +133,7 @@
 The project modules description
 """
 kwargs_common = {
-    "include_dirs": [numpy.get_include()] + _project_backend_dir,
+    "include_dirs": [numpy.get_include(), dpctl.get_include()] + _project_backend_dir,
     "extra_compile_args": _sdl_cflags,
     "extra_link_args": _project_extra_link_args,
     "define_macros": [("NPY_NO_DEPRECATED_API", "NPY_1_7_API_VERSION")],
diff --git a/tests/test_sycl_queue.py b/tests/test_sycl_queue.py

Original file line number	Diff line number	Diff line change
`@@ -63,6 +63,7 @@ __all__ = [`
`63`	`63`	`"_get_linear_index",`
`64`	`64`	`"normalize_axis",`
`65`	`65`	`"_object_to_tuple",`
	`66`	`+ "unwrap_array",`
`66`	`67`	`"use_origin_backend"`
`67`	`68`	`]`
`68`	`69`