Add support for kernel specific array views

Author: ZzEeKkAa

numba_dpex/core/kernel_interface/arrayobj.py

@@ -12,6 +12,7 @@
 
 
 from numba.core import cgutils, types
+from numba.np import arrayobj as np_arrayobj
 
 from numba_dpex.core import types as dpex_types
 
@@ -79,3 +80,58 @@ def populate_array(array, data, shape, strides, itemsize):
         setattr(array, k, v)
 
     return array
+
+
+def make_view(context, builder, aryty, ary, return_type, data, shapes, strides):
+    """
+    Build a view over the given array with the given parameters.
+
+    This is analog of numpy.np.arrayobj.make_view without parent and
+    meminfo fields, because they don't make sense on device. This function
+    intended to be used only in kernel targets.
+    """
+    retary = np_arrayobj.make_array(return_type)(context, builder)
+    populate_array(
+        retary, data=data, shape=shapes, strides=strides, itemsize=ary.itemsize
+    )
+    return retary
+
+
+def _getitem_array_generic(
+    context, builder, return_type, aryty, ary, index_types, indices
+):
+    """
+    Return the result of indexing *ary* with the given *indices*,
+    returning either a scalar or a view.
+
+    This is analog of numpy.np.arrayobj._getitem_array_generic without parent
+    and meminfo fields, because they don't make sense on device. This function
+    intended to be used only in kernel targets.
+    """
+    dataptr, view_shapes, view_strides = np_arrayobj.basic_indexing(
+        context,
+        builder,
+        aryty,
+        ary,
+        index_types,
+        indices,
+        boundscheck=context.enable_boundscheck,
+    )
+
+    if isinstance(return_type, types.Buffer):
+        # Build array view
+        retary = make_view(
+            context,
+            builder,
+            aryty,
+            ary,
+            return_type,
+            dataptr,
+            view_shapes,
+            view_strides,
+        )
+        return retary._getvalue()
+    else:
+        # Load scalar from 0-d result
+        assert not view_shapes
+        return np_arrayobj.load_item(context, builder, aryty, dataptr)

numba_dpex/dpnp_iface/arrayobj.py

@@ -6,15 +6,20 @@
 
 import dpnp
 from numba import errors, types
-from numba.core.imputils import lower_builtin
+from numba.core.imputils import impl_ret_borrowed, lower_builtin
 from numba.core.types import scalars
 from numba.core.types.containers import UniTuple
 from numba.core.typing.npydecl import parse_dtype as _ty_parse_dtype
 from numba.core.typing.npydecl import parse_shape as _ty_parse_shape
 from numba.extending import overload, overload_attribute
-from numba.np.arrayobj import getitem_arraynd_intp as np_getitem_arraynd_intp
+from numba.np.arrayobj import _getitem_array_generic as np_getitem_array_generic
+from numba.np.arrayobj import make_array
 from numba.np.numpy_support import is_nonelike
 
+from numba_dpex.core.kernel_interface.arrayobj import (
+    _getitem_array_generic as kernel_getitem_array_generic,
+)
+from numba_dpex.core.targets.kernel_target import DpexKernelTargetContext
 from numba_dpex.core.types import DpnpNdArray
 
 from ._intrinsic import (
@@ -1077,9 +1082,24 @@ def getitem_arraynd_intp(context, builder, sig, args):
     that when returning a view of a dpnp.ndarray the sycl::queue pointer
     member in the LLVM IR struct gets properly updated.
     """
-    ret = np_getitem_arraynd_intp(context, builder, sig, args)
+    getitem_call_in_kernel = isinstance(context, DpexKernelTargetContext)
+    _getitem_array_generic = np_getitem_array_generic
+
+    if getitem_call_in_kernel:
+        _getitem_array_generic = kernel_getitem_array_generic
+
+    aryty, idxty = sig.args
+    ary, idx = args
+
+    assert aryty.ndim >= 1
+    ary = make_array(aryty)(context, builder, ary)
+
+    res = _getitem_array_generic(
+        context, builder, sig.return_type, aryty, ary, (idxty,), (idx,)
+    )
+    ret = impl_ret_borrowed(context, builder, sig.return_type, res)
 
-    if isinstance(sig.return_type, DpnpNdArray):
+    if isinstance(sig.return_type, DpnpNdArray) and not getitem_call_in_kernel:
         array_val = args[0]
         array_ty = sig.args[0]
         sycl_queue_attr_pos = context.data_model_manager.lookup(