Merge pull request #1083 from IntelPython/feature/improved_sycl_queue_support

Feature/improved sycl queue support

Author: Diptorup Deb

numba_dpex/core/datamodel/models.py

@@ -3,10 +3,10 @@
 # SPDX-License-Identifier: Apache-2.0
 
 from numba.core import datamodel, types
-from numba.core.datamodel.models import ArrayModel as DpnpNdArrayModel
 from numba.core.datamodel.models import PrimitiveModel, StructModel
 from numba.core.extending import register_model
 
+from numba_dpex.core.exceptions import UnreachableError
 from numba_dpex.utils import address_space
 
 from ..types import Array, DpctlSyclQueue, DpnpNdArray, USMNdArray
@@ -23,7 +23,7 @@ def __init__(self, dmm, fe_type):
         super(GenericPointerModel, self).__init__(dmm, fe_type, be_type)
 
 
-class ArrayModel(StructModel):
+class USMArrayModel(StructModel):
     """A data model to represent a Dpex's array types in LLVM IR.
 
     Dpex's ArrayModel is based on Numba's ArrayModel for NumPy arrays. The
@@ -40,18 +40,75 @@ def __init__(self, dmm, fe_type):
             ),
             (
                 "parent",
-                types.CPointer(fe_type.dtype, addrspace=fe_type.addrspace),
+                types.CPointer(types.pyobject, addrspace=fe_type.addrspace),
             ),
             ("nitems", types.intp),
             ("itemsize", types.intp),
             (
                 "data",
                 types.CPointer(fe_type.dtype, addrspace=fe_type.addrspace),
             ),
+            (
+                "sycl_queue",
+                types.CPointer(types.void, addrspace=fe_type.addrspace),
+            ),
+            ("shape", types.UniTuple(types.intp, ndim)),
+            ("strides", types.UniTuple(types.intp, ndim)),
+        ]
+        super(USMArrayModel, self).__init__(dmm, fe_type, members)
+
+
+class DpnpNdArrayModel(StructModel):
+    """Data model for the DpnpNdArray type.
+
+    DpnpNdArrayModel is used by the numba_dpex.types.DpnpNdArray type and
+    abstracts the usmarystruct_t C type defined in
+    numba_dpex.core.runtime._usmarraystruct.h.
+
+    The DpnpNdArrayModel differs from numba's ArrayModel by including an extra
+    member sycl_queue that maps to _usmarraystruct.sycl_queue pointer. The
+    _usmarraystruct.sycl_queue pointer stores the C++ sycl::queue pointer that
+    was used to allocate the data for the dpnp.ndarray represented by an
+    instance of _usmarraystruct.
+    """
+
+    def __init__(self, dmm, fe_type):
+        ndim = fe_type.ndim
+        members = [
+            ("meminfo", types.MemInfoPointer(fe_type.dtype)),
+            ("parent", types.pyobject),
+            ("nitems", types.intp),
+            ("itemsize", types.intp),
+            ("data", types.CPointer(fe_type.dtype)),
+            ("sycl_queue", types.voidptr),
             ("shape", types.UniTuple(types.intp, ndim)),
             ("strides", types.UniTuple(types.intp, ndim)),
         ]
-        super(ArrayModel, self).__init__(dmm, fe_type, members)
+        super(DpnpNdArrayModel, self).__init__(dmm, fe_type, members)
+
+    @property
+    def flattened_field_count(self):
+        """Return the number of fields in an instance of a DpnpNdArrayModel."""
+        flattened_member_count = 0
+        members = self._members
+        for member in members:
+            if isinstance(member, types.UniTuple):
+                flattened_member_count += member.count
+            elif isinstance(
+                member,
+                (
+                    types.scalars.Integer,
+                    types.misc.PyObject,
+                    types.misc.RawPointer,
+                    types.misc.CPointer,
+                    types.misc.MemInfoPointer,
+                ),
+            ):
+                flattened_member_count += 1
+            else:
+                raise UnreachableError
+
+        return flattened_member_count
 
 
 class SyclQueueModel(StructModel):
@@ -84,7 +141,7 @@ def __init__(self, dmm, fe_type):
 def _init_data_model_manager():
     dmm = datamodel.default_manager.copy()
     dmm.register(types.CPointer, GenericPointerModel)
-    dmm.register(Array, ArrayModel)
+    dmm.register(Array, USMArrayModel)
     return dmm
 
 
@@ -103,8 +160,8 @@ def _init_data_model_manager():
 # object.
 
 # Register the USMNdArray type with the dpex ArrayModel
-register_model(USMNdArray)(ArrayModel)
-dpex_data_model_manager.register(USMNdArray, ArrayModel)
+register_model(USMNdArray)(USMArrayModel)
+dpex_data_model_manager.register(USMNdArray, USMArrayModel)
 
 # Register the DpnpNdArray type with the Numba ArrayModel
 register_model(DpnpNdArray)(DpnpNdArrayModel)

numba_dpex/core/kernel_interface/arg_pack_unpacker.py

@@ -54,6 +54,8 @@ def _unpack_usm_array(self, val):
         unpacked_array_attrs.append(ctypes.c_longlong(size))
         unpacked_array_attrs.append(ctypes.c_longlong(itemsize))
         unpacked_array_attrs.append(buf)
+        # queue: unused and passed as void*
+        unpacked_array_attrs.append(ctypes.c_size_t(0))
         for ax in range(ndim):
             unpacked_array_attrs.append(ctypes.c_longlong(shape[ax]))
         for ax in range(ndim):

numba_dpex/core/parfors/parfor_lowerer.py

@@ -5,7 +5,7 @@
 import copy
 
 from llvmlite import ir as llvmir
-from numba.core import cgutils, ir, types
+from numba.core import ir, types
 from numba.parfors.parfor import (
     find_potential_aliases_parfor,
     get_parfor_outputs,
@@ -26,6 +26,8 @@
     create_reduction_remainder_kernel_for_parfor,
 )
 
+from numba_dpex.core.datamodel.models import dpex_data_model_manager as dpex_dmm
+
 # A global list of kernels to keep the objects alive indefinitely.
 keep_alive_kernels = []
 
@@ -114,8 +116,8 @@ def _build_kernel_arglist(self, kernel_fn, lowerer):
         # kernel_fn.kernel_args as arrays get flattened.
         for arg_type in kernel_fn.kernel_arg_types:
             if isinstance(arg_type, DpnpNdArray):
-                # FIXME: Remove magic constants
-                num_flattened_args += 5 + (2 * arg_type.ndim)
+                datamodel = dpex_dmm.lookup(arg_type)
+                num_flattened_args += datamodel.flattened_field_count
             elif arg_type == types.complex64 or arg_type == types.complex128:
                 num_flattened_args += 2
             else:
@@ -134,15 +136,16 @@ def _build_kernel_arglist(self, kernel_fn, lowerer):
             argtype = kernel_fn.kernel_arg_types[arg_num]
             llvm_val = _getvar(lowerer, arg)
             if isinstance(argtype, DpnpNdArray):
+                datamodel = dpex_dmm.lookup(arg_type)
                 self.kernel_builder.build_array_arg(
                     array_val=llvm_val,
+                    array_data_model=datamodel,
                     array_rank=argtype.ndim,
                     arg_list=self.args_list,
                     args_ty_list=self.args_ty_list,
                     arg_num=self.kernel_arg_num,
                 )
-                # FIXME: Get rid of magic constants
-                self.kernel_arg_num += 5 + (2 * argtype.ndim)
+                self.kernel_arg_num += datamodel.flattened_field_count
             else:
                 if argtype == types.complex64:
                     self.kernel_builder.build_complex_arg(

numba_dpex/core/runtime/_dpexrt_python.c

@@ -21,7 +21,7 @@
 
 #include "_dbg_printer.h"
 #include "_queuestruct.h"
-#include "numba/_arraystruct.h"
+#include "_usmarraystruct.h"
 
 // forward declarations
 static struct PyUSMArrayObject *PyUSMNdArray_ARRAYOBJ(PyObject *obj);
@@ -49,14 +49,14 @@ static NRT_MemInfo *DPEXRT_MemInfo_alloc(npy_intp size,
                                          size_t usm_type,
                                          const DPCTLSyclQueueRef qref);
 static void usmndarray_meminfo_dtor(void *ptr, size_t size, void *info);
-static PyObject *box_from_arystruct_parent(arystruct_t *arystruct,
+static PyObject *box_from_arystruct_parent(usmarystruct_t *arystruct,
                                            int ndim,
                                            PyArray_Descr *descr);
 
 static int DPEXRT_sycl_usm_ndarray_from_python(PyObject *obj,
-                                               arystruct_t *arystruct);
+                                               usmarystruct_t *arystruct);
 static PyObject *
-DPEXRT_sycl_usm_ndarray_to_python_acqref(arystruct_t *arystruct,
+DPEXRT_sycl_usm_ndarray_to_python_acqref(usmarystruct_t *arystruct,
                                          PyTypeObject *retty,
                                          int ndim,
                                          int writeable,
@@ -770,7 +770,7 @@ static npy_intp product_of_shape(npy_intp *shape, npy_intp ndim)
  * @return   {return}       Error code representing success (0) or failure (-1).
  */
 static int DPEXRT_sycl_usm_ndarray_from_python(PyObject *obj,
-                                               arystruct_t *arystruct)
+                                               usmarystruct_t *arystruct)
 {
     struct PyUSMArrayObject *arrayobj = NULL;
     int i = 0, j = 0, k = 0, ndim = 0, exp = 0;
@@ -827,6 +827,7 @@ static int DPEXRT_sycl_usm_ndarray_from_python(PyObject *obj,
     }
 
     arystruct->data = data;
+    arystruct->sycl_queue = qref;
     arystruct->nitems = nitems;
     arystruct->itemsize = itemsize;
     arystruct->parent = obj;
@@ -892,7 +893,7 @@ static int DPEXRT_sycl_usm_ndarray_from_python(PyObject *obj,
  * @return   {return}       A PyObject created from the arystruct_t->parent, if
  *                          the PyObject could not be created return NULL.
  */
-static PyObject *box_from_arystruct_parent(arystruct_t *arystruct,
+static PyObject *box_from_arystruct_parent(usmarystruct_t *arystruct,
                                            int ndim,
                                            PyArray_Descr *descr)
 {
@@ -914,8 +915,10 @@ static PyObject *box_from_arystruct_parent(arystruct_t *arystruct,
     }
 
     if ((void *)UsmNDArray_GetData(arrayobj) != arystruct->data) {
-        DPEXRT_DEBUG(drt_debug_print("DPEXRT-DEBUG: Arrayobj cannot be boxed "
-                                     "from parent as data pointer is NULL.\n"));
+        DPEXRT_DEBUG(drt_debug_print(
+            "DPEXRT-DEBUG: Arrayobj cannot be boxed "
+            "from parent as data pointer in the arystruct is not the same as "
+            "the data pointer in the parent object.\n"));
         return NULL;
     }
 
@@ -978,7 +981,7 @@ static PyObject *box_from_arystruct_parent(arystruct_t *arystruct,
  *
  */
 static PyObject *
-DPEXRT_sycl_usm_ndarray_to_python_acqref(arystruct_t *arystruct,
+DPEXRT_sycl_usm_ndarray_to_python_acqref(usmarystruct_t *arystruct,
                                          PyTypeObject *retty,
                                          int ndim,
                                          int writeable,
@@ -1094,8 +1097,7 @@ DPEXRT_sycl_usm_ndarray_to_python_acqref(arystruct_t *arystruct,
     typenum = descr->type_num;
     usm_ndarr_obj = UsmNDArray_MakeFromPtr(
         ndim, shape, typenum, strides, (DPCTLSyclUSMRef)arystruct->data,
-        (DPCTLSyclQueueRef)miobj->meminfo->external_allocator->opaque_data, 0,
-        (PyObject *)miobj);
+        (DPCTLSyclQueueRef)arystruct->sycl_queue, 0, (PyObject *)miobj);
 
     if (usm_ndarr_obj == NULL ||
         !PyObject_TypeCheck(usm_ndarr_obj, &PyUSMArrayType))

numba_dpex/core/runtime/_queuestruct.h

@@ -1,10 +1,15 @@
-#ifndef NUMBA_DPEX_QUEUESTRUCT_H_
-#define NUMBA_DPEX_QUEUESTRUCT_H_
-/*
- * Fill in the *queuestruct* with information from the Numpy array *obj*.
- * *queuestruct*'s layout is defined in numba.targets.arrayobj (look
- * for the ArrayTemplate class).
- */
+// SPDX-FileCopyrightText: 2020 - 2023 Intel Corporation
+//
+// SPDX-License-Identifier: Apache-2.0
+
+//===----------------------------------------------------------------------===//
+///
+/// \file
+/// Defines the numba-dpex native representation for a dpctl.SyclQueue
+///
+//===----------------------------------------------------------------------===//
+
+#pragma once
 
 #include <Python.h>
 
@@ -13,5 +18,3 @@ typedef struct
     PyObject *parent;
     void *queue_ref;
 } queuestruct_t;
-
-#endif /* NUMBA_DPEX_QUEUESTRUCT_H_ */

numba_dpex/core/runtime/_usmarraystruct.h

@@ -0,0 +1,27 @@
+// SPDX-FileCopyrightText: 2020 - 2023 Intel Corporation
+//
+// SPDX-License-Identifier: Apache-2.0
+
+//===----------------------------------------------------------------------===//
+///
+/// \file
+/// Defines the numba-dpex native representation for a dpctl.tensor.usm_ndarray
+///
+//===----------------------------------------------------------------------===//
+
+#pragma once
+
+#include <Python.h>
+#include <numpy/npy_common.h>
+
+typedef struct
+{
+    void *meminfo;
+    PyObject *parent;
+    npy_intp nitems;
+    npy_intp itemsize;
+    void *data;
+    void *sycl_queue;
+
+    npy_intp shape_and_strides[];
+} usmarystruct_t;