Use different data models for DpnpNdArray Type for kernel and dpjit targets

numba_dpex/core/datamodel/models.py

@@ -57,6 +57,30 @@ def __init__(self, dmm, fe_type):
         ]
         super(USMArrayModel, self).__init__(dmm, fe_type, members)
 
+    @property
+    def flattened_field_count(self):
+        """Return the number of fields in an instance of a USMArrayModel."""
+        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 DpnpNdArrayModel(StructModel):
     """Data model for the DpnpNdArray type.
@@ -138,35 +162,54 @@ def __init__(self, dmm, fe_type):
         super(SyclQueueModel, self).__init__(dmm, fe_type, members)
 
 
-def _init_data_model_manager():
+def _init_data_model_manager() -> datamodel.DataModelManager:
+    """Initializes a DpexKernelTarget-specific data model manager.
+
+    SPIRV kernel functions for certain types of devices require an explicit
+    address space qualifier for pointers. For OpenCL HD Graphics
+    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.
+
+    Returns:
+        DataModelManager: A numba-dpex DpexKernelTarget-specific data model
+        manager
+    """
     dmm = datamodel.default_manager.copy()
     dmm.register(types.CPointer, GenericPointerModel)
     dmm.register(Array, USMArrayModel)
+
+    # 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 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 DpctlSyclQueue type to SyclQueueModel in numba_dpex's data
+    # model manager. The dpex_data_model_manager is used by the DpexKernelTarget
+    dmm.register(DpctlSyclQueue, SyclQueueModel)
+
     return dmm
 
 
 dpex_data_model_manager = _init_data_model_manager()
 
-# XXX A kernel function has the spir_kernel ABI and requires pointers to have an
-# address space attribute. For this reason, the UsmNdArray type uses dpex's
-# ArrayModel where the pointers are address space casted to have a SYCL-specific
-# address space value. The DpnpNdArray type can be used inside djit functions
-# as host function calls arguments, such as dpnp library calls. The DpnpNdArray
-# needs to use Numba's array model as its data model. Thus, from a Numba typing
-# perspective dpnp.ndarrays cannot be directly passed to a kernel. To get
-# around the limitation, the DpexKernelTypingContext does not resolve the type
-# of dpnp.array args to a kernel as DpnpNdArray type objects, but uses the
-# ``to_usm_ndarray`` utility function to convert them into a UsmNdArray type
-# object.
-
-# Register the USMNdArray type with the dpex ArrayModel
+
+# Register the USMNdArray type to USMArrayModel in numba's default data model
+# manager
 register_model(USMNdArray)(USMArrayModel)
-dpex_data_model_manager.register(USMNdArray, USMArrayModel)
 
-# Register the DpnpNdArray type with the Numba ArrayModel
+# Register the DpnpNdArray type to DpnpNdArrayModel in numba's default data
+# model manager
 register_model(DpnpNdArray)(DpnpNdArrayModel)
-dpex_data_model_manager.register(DpnpNdArray, DpnpNdArrayModel)
 
 # Register the DpctlSyclQueue type
 register_model(DpctlSyclQueue)(SyclQueueModel)
-dpex_data_model_manager.register(DpctlSyclQueue, SyclQueueModel)

numba_dpex/core/parfors/kernel_builder.py

@@ -70,30 +70,6 @@ def _compile_kernel_parfor(
         func_ir, kernel_name
     )
 
-    # A cast from DpnpNdArray type to USMNdArray is needed for all arguments of
-    # DpnpNdArray type. Although, DpnpNdArray derives from USMNdArray the two
-    # types use different data models. USMNdArray uses the
-    # numba_dpex.core.datamodel.models.ArrayModel data model that defines all
-    # CPointer type members in the GLOBAL address space. The DpnpNdArray uses
-    # Numba's default ArrayModel that does not define pointers in any specific
-    # address space. For OpenCL HD Graphics devices, defining a kernel function
-    # (spir_kernel calling convention) with pointer arguments that have no
-    # address space qualifier causes a run time crash. By casting the argument
-    # type for parfor arguments from DpnpNdArray type to the USMNdArray type the
-    # generated kernel always has an address space qualifier, avoiding the issue
-    # on OpenCL HD graphics devices.
-
-    for i, argty in enumerate(argtypes):
-        if isinstance(argty, DpnpNdArray):
-            new_argty = USMNdArray(
-                ndim=argty.ndim,
-                layout=argty.layout,
-                dtype=argty.dtype,
-                usm_type=argty.usm_type,
-                queue=argty.queue,
-            )
-            argtypes[i] = new_argty
-
     # compile the kernel
     kernel.compile(
         args=argtypes,

numba_dpex/core/targets/kernel_target.py

@@ -72,28 +72,6 @@ def resolve_argument_type(self, val):
                     type=str(type(val)), value=val
                 )
 
-            # A cast from DpnpNdArray type to USMNdArray is needed for all
-            # arguments of DpnpNdArray type. Although, DpnpNdArray derives from
-            # USMNdArray the two types use different data models. USMNdArray
-            # uses the numba_dpex.core.datamodel.models.ArrayModel data model
-            # that defines all CPointer type members in the GLOBAL address
-            # space. The DpnpNdArray uses Numba's default ArrayModel that does
-            # not define pointers in any specific address space. For OpenCL HD
-            # Graphics devices, defining a kernel function (spir_kernel calling
-            # convention) with pointer arguments that have no address space
-            # qualifier causes a run time crash. By casting the argument type
-            # for parfor arguments from DpnpNdArray type to the USMNdArray type
-            # the generated kernel always has an address space qualifier,
-            # avoiding the issue on OpenCL HD graphics devices.
-            if isinstance(numba_type, DpnpNdArray):
-                return USMNdArray(
-                    ndim=numba_type.ndim,
-                    layout=numba_type.layout,
-                    dtype=numba_type.dtype,
-                    usm_type=numba_type.usm_type,
-                    queue=numba_type.queue,
-                )
-
         except ValueError:
             # When an array-like kernel argument is not recognized by
             # numba-dpex, this additional check sees if the array-like object

numba_dpex/dpnp_iface/arrayobj.py

@@ -1075,12 +1075,13 @@ def getitem_arraynd_intp(context, builder, sig, args):
     """
     ret = np_getitem_arraynd_intp(context, builder, sig, args)
 
-    array_val = args[0]
-    array_ty = sig.args[0]
-    sycl_queue_attr_pos = dpex_dmm.lookup(array_ty).get_field_position(
-        "sycl_queue"
-    )
-    sycl_queue_attr = builder.extract_value(array_val, sycl_queue_attr_pos)
-    ret = builder.insert_value(ret, sycl_queue_attr, sycl_queue_attr_pos)
+    if isinstance(sig.return_type, DpnpNdArray):
+        array_val = args[0]
+        array_ty = sig.args[0]
+        sycl_queue_attr_pos = dpex_dmm.lookup(array_ty).get_field_position(
+            "sycl_queue"
+        )
+        sycl_queue_attr = builder.extract_value(array_val, sycl_queue_attr_pos)
+        ret = builder.insert_value(ret, sycl_queue_attr, sycl_queue_attr_pos)
 
     return ret

numba_dpex/tests/core/types/DpnpNdArray/test_models.py

@@ -3,26 +3,25 @@
 # SPDX-License-Identifier: Apache-2.0
 
 from numba import types
-from numba.core.datamodel import models
+from numba.core.datamodel import default_manager, models
 
 from numba_dpex.core.datamodel.models import (
     DpnpNdArrayModel,
+    USMArrayModel,
     dpex_data_model_manager,
 )
 from numba_dpex.core.types.dpnp_ndarray_type import DpnpNdArray
 
 
 def test_model_for_DpnpNdArray():
-    """Test that model is registered for DpnpNdArray instances.
-
-    The model for DpnpNdArray is dpex's ArrayModel.
-
+    """Test the datamodel for DpnpNdArray that is registered with numba's
+    default datamodel manager and numba_dpex's kernel data model manager.
     """
-
-    model = dpex_data_model_manager.lookup(
-        DpnpNdArray(ndim=1, dtype=types.float64, layout="C")
-    )
-    assert isinstance(model, DpnpNdArrayModel)
+    dpnp_ndarray = DpnpNdArray(ndim=1, dtype=types.float64, layout="C")
+    model = dpex_data_model_manager.lookup(dpnp_ndarray)
+    assert isinstance(model, USMArrayModel)
+    default_model = default_manager.lookup(dpnp_ndarray)
+    assert isinstance(default_model, DpnpNdArrayModel)
 
 
 def test_dpnp_ndarray_Model():