Merge pull request #1263 from IntelPython/feature/add_dpctl_tensor_su…

…pport_to_dpjit

Add dpctl tensor support

numba_dpex/core/datamodel/models.py

@@ -55,58 +55,68 @@ def __init__(self, dmm, fe_type):
         super(GenericPointerModel, self).__init__(dmm, fe_type, be_type)
 
 
-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
-    dpex model adds an extra address space attribute to all pointer members
-    in the array.
+class USMArrayDeviceModel(StructModel):
+    """A data model to represent a usm array type in the LLVM IR generated for a
+    device-only kernel function.
+
+    The USMArrayDeviceModel adds an extra address space attribute to the data
+    member. The extra attribute is needed when passing usm_ndarray array
+    arguments to kernels that are compiled for certain OpenCL GPU devices. Note
+    that the address space attribute is applied only to the data member and not
+    other members of USMArrayDeviceModel that are pointers. It is done this way
+    as other pointer members such as meminfo are not used inside a kernel and
+    these members maybe removed from the USMArrayDeviceModel in
+    future (refer #929).
+
+    We use separate data models for host (USMArrayHostModel) and device
+    (USMArrayDeviceModel) as the address space attribute is only required for
+    kernel functions and not needed for functions that are compiled for a host
+    memory space.
     """
 
+    # TODO: Evaluate the need to pass meminfo and parent attributes of an array
+    #   as kernel params: https://github.com/IntelPython/numba-dpex/issues/929
+
     def __init__(self, dmm, fe_type):
         ndim = fe_type.ndim
         members = [
-            (
-                "meminfo",
-                types.CPointer(fe_type.dtype, addrspace=fe_type.addrspace),
-            ),
-            (
-                "parent",
-                types.CPointer(types.pyobject, addrspace=fe_type.addrspace),
-            ),
+            # meminfo never used in kernel, so we don'te care about addrspace
+            ("meminfo", types.MemInfoPointer(fe_type.dtype)),
+            # parent never used in kernel, so we don'te care about addrspace
+            ("parent", types.pyobject),
             ("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),
-            ),
+            # sycl_queue never used in kernel, so we don'te care about addrspace
+            ("sycl_queue", types.voidptr),
             ("shape", types.UniTuple(types.intp, ndim)),
             ("strides", types.UniTuple(types.intp, ndim)),
         ]
-        super(USMArrayModel, self).__init__(dmm, fe_type, members)
+        super(USMArrayDeviceModel, self).__init__(dmm, fe_type, members)
 
     @property
     def flattened_field_count(self):
-        """Return the number of fields in an instance of a USMArrayModel."""
+        """
+        Return the number of fields in an instance of a USMArrayDeviceModel.
+        """
         return _get_flattened_member_count(self)
 
 
-class DpnpNdArrayModel(StructModel):
-    """Data model for the DpnpNdArray type.
+class USMArrayHostModel(StructModel):
+    """Data model for the USMNdArray type when used in a host-only function.
 
-    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.
+    USMArrayHostModel is used by the numba_dpex.types.USMNdArray and
+    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
+    The USMArrayDeviceModel 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.
+    was used to allocate the data for the dpctl.tensor.usm_ndarray or
+    dpnp.ndarray represented by an instance of _usmarraystruct.
     """
 
     def __init__(self, dmm, fe_type):
@@ -121,11 +131,11 @@ def __init__(self, dmm, fe_type):
             ("shape", types.UniTuple(types.intp, ndim)),
             ("strides", types.UniTuple(types.intp, ndim)),
         ]
-        super(DpnpNdArrayModel, self).__init__(dmm, fe_type, members)
+        super(USMArrayHostModel, self).__init__(dmm, fe_type, members)
 
     @property
     def flattened_field_count(self):
-        """Return the number of fields in an instance of a DpnpNdArrayModel."""
+        """Return the number of fields in an instance of a USMArrayHostModel."""
         return _get_flattened_member_count(self)
 
 
@@ -242,28 +252,28 @@ def _init_data_model_manager() -> datamodel.DataModelManager:
     devices, defining a kernel function (spir_kernel calling convention) with
     pointer arguments that have no address space qualifier causes a run time
     crash. For this reason, numba-dpex defines two separate data
-    models: USMArrayModel and DpnpNdArrayModel. When a dpnp.ndarray object is
-    passed as an argument to a ``numba_dpex.kernel`` decorated function it uses
-    the USMArrayModel and when passed to a ``numba_dpex.dpjit`` decorated
-    function it uses the DpnpNdArrayModel. The difference is due to the fact
-    that inside a ``dpjit`` decorated function a dpnp.ndarray object can be
-    passed to any other regular function.
+    models: USMArrayDeviceModel and USMArrayHostModel. When a dpnp.ndarray
+    object is passed as an argument to a ``numba_dpex.kernel`` decorated
+    function it uses the USMArrayDeviceModel and when passed to a
+    ``numba_dpex.dpjit`` decorated function it uses the USMArrayHostModel.
+    The difference is due to the fact that inside a ``dpjit`` decorated function
+    a dpnp.ndarray object can be passed to any other regular function.
 
     Returns:
         DataModelManager: A numba-dpex DpexKernelTarget-specific data model
         manager
     """
     dmm = datamodel.default_manager.copy()
     dmm.register(types.CPointer, GenericPointerModel)
-    dmm.register(Array, USMArrayModel)
+    dmm.register(Array, USMArrayDeviceModel)
 
-    # Register the USMNdArray type to USMArrayModel in numba_dpex's data model
-    # manager. The dpex_data_model_manager is used by the DpexKernelTarget
-    dmm.register(USMNdArray, USMArrayModel)
+    # Register the USMNdArray type to USMArrayDeviceModel in numba_dpex's data
+    # model manager. The dpex_data_model_manager is used by the DpexKernelTarget
+    dmm.register(USMNdArray, USMArrayDeviceModel)
 
-    # Register the DpnpNdArray type to USMArrayModel in numba_dpex's data model
-    # manager. The dpex_data_model_manager is used by the DpexKernelTarget
-    dmm.register(DpnpNdArray, USMArrayModel)
+    # Register the DpnpNdArray type to USMArrayDeviceModel in numba_dpex's data
+    # model manager. The dpex_data_model_manager is used by the DpexKernelTarget
+    dmm.register(DpnpNdArray, USMArrayDeviceModel)
 
     # Register the DpctlSyclQueue type to SyclQueueModel in numba_dpex's data
     # model manager. The dpex_data_model_manager is used by the DpexKernelTarget
@@ -275,13 +285,13 @@ def _init_data_model_manager() -> datamodel.DataModelManager:
 dpex_data_model_manager = _init_data_model_manager()
 
 
-# Register the USMNdArray type to USMArrayModel in numba's default data model
-# manager
-register_model(USMNdArray)(USMArrayModel)
+# Register the USMNdArray type to USMArrayDeviceModel in numba's default data
+# model manager
+register_model(USMNdArray)(USMArrayHostModel)
 
-# Register the DpnpNdArray type to DpnpNdArrayModel in numba's default data
+# Register the DpnpNdArray type to USMArrayHostModel in numba's default data
 # model manager
-register_model(DpnpNdArray)(DpnpNdArrayModel)
+register_model(DpnpNdArray)(USMArrayHostModel)
 
 # Register the DpctlSyclQueue type
 register_model(DpctlSyclQueue)(SyclQueueModel)

numba_dpex/core/runtime/_dpexrt_python.c

@@ -740,12 +740,19 @@ static struct PyUSMArrayObject *PyUSMNdArray_ARRAYOBJ(PyObject *obj)
 {
     PyObject *arrayobj = NULL;
 
-    arrayobj = PyObject_GetAttrString(obj, "_array_obj");
+    if (PyObject_TypeCheck(obj, &PyUSMArrayType)) {
+        DPEXRT_DEBUG(
+            drt_debug_print("DPEXRT-DEBUG: usm array was passed directly\n"));
+        arrayobj = obj;
+    }
+    else if (PyObject_HasAttrString(obj, "_array_obj")) {
+        arrayobj = PyObject_GetAttrString(obj, "_array_obj");
 
-    if (!arrayobj)
-        return NULL;
-    if (!PyObject_TypeCheck(arrayobj, &PyUSMArrayType))
-        return NULL;
+        if (!arrayobj)
+            return NULL;
+        if (!PyObject_TypeCheck(arrayobj, &PyUSMArrayType))
+            return NULL;
+    }
 
     struct PyUSMArrayObject *pyusmarrayobj =
         (struct PyUSMArrayObject *)(arrayobj);
@@ -1164,6 +1171,8 @@ DPEXRT_sycl_usm_ndarray_to_python_acqref(usmarystruct_t *arystruct,
         return MOD_ERROR_VAL;
     }
 
+    // TODO: check if the object is dpctl tensor and return usm_ndarr_obj then
+
     //  call new on dpnp_array
     dpnp_array_mod = PyImport_ImportModule("dpnp.dpnp_array");
     if (!dpnp_array_mod) {

numba_dpex/core/types/dpnp_ndarray_type.py

@@ -200,19 +200,22 @@ def __allocate__(
         return out
 
 
+# TODO: move this section to separate file
 # --------------- Boxing/Unboxing logic for dpnp.ndarray ----------------------#
 
 
-@unbox(DpnpNdArray)
+@unbox(USMNdArray)
 def unbox_dpnp_nd_array(typ, obj, c):
-    """Converts a dpnp.ndarray object to a Numba internal array structure.
+    """Converts a dpctl.tensor.usm_ndarray/dpnp.ndarray object to a Numba-dpex
+    internal array structure.
 
     Args:
         typ : The Numba type of the PyObject
         obj : The actual PyObject to be unboxed
         c : The unboxing context
 
-    Returns: A NativeValue object representing an unboxed dpnp.ndarray
+    Returns: A NativeValue object representing an unboxed
+        dpctl.tensor.usm_ndarray/dpnp.ndarray
     """
     # Reusing the numba.core.base.BaseContext's make_array function to get a
     # struct allocated. The same struct is used for numpy.ndarray
@@ -264,24 +267,26 @@ def unbox_dpnp_nd_array(typ, obj, c):
     with c.builder.if_then(failed, likely=False):
         c.pyapi.err_set_string(
             "PyExc_TypeError",
-            "can't unbox array from PyObject into "
+            "can't unbox usm array from PyObject into "
             "native value.  The object maybe of a "
             "different type",
         )
     return NativeValue(c.builder.load(aryptr), is_error=failed)
 
 
-@box(DpnpNdArray)
+@box(USMNdArray)
 def box_array(typ, val, c):
-    """Boxes a NativeValue representation of DpnpNdArray type into a
-    dpnp.ndarray PyObject
+    """Boxes a NativeValue representation of USMNdArray/DpnpNdArray type into a
+    dpctl.tensor.usm_ndarray/dpnp.ndarray PyObject
 
     Args:
-        typ: The representation of the DpnpNdArray type.
-        val: A native representation of a Numba DpnpNdArray type object.
+        typ: The representation of the USMNdArray/DpnpNdArray type.
+        val: A native representation of a Numba USMNdArray/DpnpNdArray type
+            object.
         c: The boxing context.
 
-    Returns: A Pyobject for a dpnp.ndarray boxed from the Numba native value.
+    Returns: A Pyobject for a dpctl.tensor.usm_ndarray/dpnp.ndarray boxed from
+        the Numba-dpex native value.
     """
     if c.context.enable_nrt:
         np_dtype = numpy_support.as_dtype(typ.dtype)

numba_dpex/core/utils/kernel_launcher.py

@@ -19,7 +19,7 @@
 from numba_dpex import config, utils
 from numba_dpex.core.exceptions import UnreachableError
 from numba_dpex.core.runtime.context import DpexRTContext
-from numba_dpex.core.types import DpnpNdArray
+from numba_dpex.core.types import USMNdArray
 from numba_dpex.core.types.range_types import NdRangeType, RangeType
 from numba_dpex.dpctl_iface import libsyclinterface_bindings as sycl
 from numba_dpex.dpctl_iface._helpers import numba_type_to_dpctl_typenum
@@ -258,7 +258,7 @@ def _build_array_arg(  # pylint: disable=too-many-arguments
         args_ty_list,
         arg_num,
     ):
-        """Creates a list of LLVM Values for an unpacked DpnpNdArray kernel
+        """Creates a list of LLVM Values for an unpacked USMNdArray kernel
         argument.
 
         The steps performed here are the same as in
@@ -524,7 +524,7 @@ def set_queue(self, sycl_queue_ref: llvmir.Instruction):
     def set_queue_from_arguments(
         self,
     ):
-        """Sets the sycl queue from the first DpnpNdArray argument provided
+        """Sets the sycl queue from the first USMNdArray argument provided
         earlier."""
         queue_ref = get_queue_from_llvm_values(
             self.context,
@@ -786,7 +786,7 @@ def _get_num_flattened_kernel_args(
         flattens dpnp arrays and complex values."""
         num_flattened_kernel_args = 0
         for arg_type in kernel_argtys:
-            if isinstance(arg_type, DpnpNdArray):
+            if isinstance(arg_type, USMNdArray):
                 datamodel = self.kernel_dmm.lookup(arg_type)
                 num_flattened_kernel_args += datamodel.flattened_field_count
             elif arg_type in [types.complex64, types.complex128]:
@@ -806,7 +806,7 @@ def _populate_kernel_args_and_args_ty_arrays(
         kernel_arg_num = 0
         for arg_num, argtype in enumerate(kernel_argtys):
             llvm_val = callargs_ptrs[arg_num]
-            if isinstance(argtype, DpnpNdArray):
+            if isinstance(argtype, USMNdArray):
                 datamodel = self.kernel_dmm.lookup(argtype)
                 self._build_array_arg(
                     array_val=llvm_val,
@@ -853,13 +853,13 @@ def get_queue_from_llvm_values(
     ll_kernel_args: list[llvmir.Instruction],
 ):
     """
-    Get the sycl queue from the first DpnpNdArray argument. Prior passes
+    Get the sycl queue from the first USMNdArray argument. Prior passes
     before lowering make sure that compute-follows-data is enforceable
     for a specific call to a kernel. As such, at the stage of lowering
-    the queue from the first DpnpNdArray argument can be extracted.
+    the queue from the first USMNdArray argument can be extracted.
     """
     for arg_num, argty in enumerate(ty_kernel_args):
-        if isinstance(argty, DpnpNdArray):
+        if isinstance(argty, USMNdArray):
             llvm_val = ll_kernel_args[arg_num]
             datamodel = ctx.data_model_manager.lookup(argty)
             sycl_queue_attr_pos = datamodel.get_field_position("sycl_queue")

numba_dpex/experimental/kernel_dispatcher.py

@@ -38,7 +38,7 @@
     CompilationMode,
     DpexKernelTargetContext,
 )
-from numba_dpex.core.types import DpnpNdArray, USMNdArray
+from numba_dpex.core.types import USMNdArray
 from numba_dpex.core.utils import kernel_launcher as kl
 
 from .target import DPEX_KERNEL_EXP_TARGET_NAME, dpex_exp_kernel_target
@@ -56,7 +56,7 @@ class _KernelCompiler(_FunctionCompiler):
     def check_queue_equivalence_of_args(
         self, py_func_name: str, args: [types.Type, ...]
     ):
-        """Evaluates if all DpnpNdArray arguments passed to a kernel function
+        """Evaluates if all USMNdArray arguments passed to a kernel function
         has the same DpctlSyclQueue type.
 
         Args:
@@ -65,15 +65,15 @@ def check_queue_equivalence_of_args(
             argument passed to the kernel
 
         Raises:
-            ExecutionQueueInferenceError: If all DpnpNdArray were not allocated
+            ExecutionQueueInferenceError: If all USMNdArray were not allocated
             on the same dpctl.SyclQueue
-            ExecutionQueueInferenceError: If there were not DpnpNdArray
+            ExecutionQueueInferenceError: If there were not USMNdArray
             arguments passed to the kernel.
         """
         common_queue = None
 
         for arg in args:
-            if isinstance(arg, DpnpNdArray):
+            if isinstance(arg, USMNdArray):
                 if common_queue is None:
                     common_queue = arg.queue
                 elif common_queue != arg.queue:
@@ -143,9 +143,9 @@ def check_arguments(self, py_func_name: str, args: [types.Type, ...]):
             KernelHasReturnValueError: non void return type.
             InvalidKernelSpecializationError: unsupported arguments where
                 provided.
-            ExecutionQueueInferenceError: If all DpnpNdArray were not allocated
+            ExecutionQueueInferenceError: If all USMNdArray were not allocated
                 on the same dpctl.SyclQueue
-            ExecutionQueueInferenceError: If there were not DpnpNdArray
+            ExecutionQueueInferenceError: If there were not USMNdArray
                 arguments passed to the kernel.
         """
         self.check_sig_types(py_func_name, args, None)
@@ -343,7 +343,7 @@ def typeof_pyval(self, val):
         # can save a couple µs.
         try:
             tp = typeof(val, Purpose.argument)
-            if isinstance(tp, types.Array) and not isinstance(tp, DpnpNdArray):
+            if isinstance(tp, types.Array) and not isinstance(tp, USMNdArray):
                 raise UnsupportedKernelArgumentError(
                     type=str(type(val)), value=val
                 )