Creating slices over dpctl.tensor.usm_ndarray objects inside a kernel cause failure

[diptorupd]

The following reproducer shows a bug in numba-dpex kernels when the input argument is a dpctl.tensor.usm_ndarray.

import numba_dpex as dpex
import numba_dpex.kernel_api as kapi
import dpnp
import dpctl.tensor as dpt


@dpex.kernel
def test_slice(item: kapi.Item, t):
    i = item.get_id(0)
    tt = t[i]


t = dpt.zeros((10, 3))
tdpnp = dpnp.zeros((10, 3))

dpex.call_kernel(test_slice, kapi.Range(10), tdpnp)
print("dpnp works....!")
dpex.call_kernel(test_slice, kapi.Range(10), t)

Creating a slice from a dpctl.tensor.usm_ndarray produces the following error:

Traceback (most recent call last):
  File "/localdisk/work/diptorup/devel/numba-dpex/slice_bug.py", line 19, in <module>
    dpex_exp.call_kernel(test_slice, kapi.Range(10), t)
  File "/localdisk/work/diptorup/devel/miniconda3/envs/dpep-devel/lib/python3.11/site-packages/numba/core/dispatcher.py", line 468, in _compile_for_args
    error_rewrite(e, 'typing')
  File "/localdisk/work/diptorup/devel/miniconda3/envs/dpep-devel/lib/python3.11/site-packages/numba/core/dispatcher.py", line 409, in error_rewrite
    raise e.with_traceback(None)
numba.core.errors.TypingError: Failed in dpex_dpjit_nopython mode pipeline (step: nopython frontend)
No implementation of function Function(<intrinsic _submit_kernel_sync>) found for signature:
 
 >>> _submit_kernel_sync(type(SPIRVKernelDispatcher(<function test_slice at 0x7f7f848d2980>)), Range<1>, UniTuple(USMNdArray(dtype=float64, ndim=2, layout=C, address_space=1, usm_type=device, device=level_zero:gpu:0, sycl_queue=DpctlSyclQueue on level_zero:gpu:0) x 1))
 
There are 2 candidate implementations:
  - Of which 2 did not match due to:
  Intrinsic in function '_submit_kernel_sync': File: ../../../../../../numba-dpex/numba_dpex/experimental/launcher.py: Line 78.
    With argument(s): '(type(SPIRVKernelDispatcher(<function test_slice at 0x7f7f848d2980>)), Range<1>, UniTuple(USMNdArray(dtype=float64, ndim=2, layout=C, address_space=1, usm_type=device, device=level_zero:gpu:0, sycl_queue=DpctlSyclQueue on level_zero:gpu:0) x 1))':
   Rejected as the implementation raised a specific error:
     LoweringError: Failed in dpex_kernel_nopython mode pipeline (step: numba_dpex qualified name disambiguation)
   "USMArrayDeviceModel does not have a field named 'meminfo'"
   
   File "slice_bug.py", line 11:
   def test_slice(item: kapi.Item, t):
       <source elided>
       i = item.get_id(0)
       tt = t[i]
       ^
   
   During: lowering "tt = getitem(value=t, index=i, fn=<built-in function getitem>)" at /localdisk/work/diptorup/devel/numba-dpex/slice_bug.py (11)
  raised from /localdisk/work/diptorup/devel/miniconda3/envs/dpep-devel/lib/python3.11/site-packages/numba/core/errors.py:874

During: resolving callee type: Function(<intrinsic _submit_kernel_sync>)
During: typing of call at /localdisk/work/diptorup/devel/numba-dpex/numba_dpex/experimental/launcher.py (237)


File "numba_dpex/experimental/launcher.py", line 237:
def call_kernel(kernel_fn, index_space, *kernel_args) -> None:
    <source elided>
    """
    _submit_kernel_sync(  # pylint: disable=E1120
    ^

The kernel works as expected when a dpnp ndarray is passed as the input. I discovered the potential bug when updating an oneapi source samples for numba-dpex: https://github.com/oneapi-src/oneAPI-samples/blob/release/2024.1/AI-and-Analytics/Features-and-Functionality/IntelPython_Numpy_Numba_dpex_kNN/IntelPython_Numpy_Numba_dpex_kNN.py

[diptorupd]

The difference between the two cases is that dpctl.tensor.usm_ndarray resolves to USMNdArray and dpnp.ndarray resolves to DpnpNdArray.

For some reason, USMNdArray type is not properly registered with USMArrayDeviceModel.

The issue is not with the data model registration. dpctl.tensor.usm_ndarray is also using the USMArrayDeviceModel. The problem is the overridden _getitem_array_generic in kernel_interface.arrayobj is used only for the DpnpNdArray:

numba-dpex/numba_dpex/dpnp_iface/arrayobj.py

Line 1073 in f6a79b4

@lower_builtin(operator.getitem, DpnpNdArray, types.Integer)

To fix for dpctl.tensor.usm_ndarray, the getitem_arraynd_intp that uses kernel_interface.arrayobj should also be registered for the USMNdArray type.