`dpt.roll` does not work with unsigned shift value

[antonwolfy]

Below code results into an unexpected exception inside _copy_usm_ndarray_for_roll_nd call:

import numpy, dpctl, dpctl.tensor as dpt

dpctl.__version__
# Out: '0.18.0+0.g786365ee72.dirty'

a = dpt.arange(10)
dpt.roll(a, numpy.uint64(1), axis=0)
---------------------------------------------------------------------------
TypeError                                 Traceback (most recent call last)
Cell In[26], line 1
----> 1 dpt.roll(a, numpy.uint64(1), axis=0)

File ~/miniforge3/envs/dpnp_dev/lib/python3.11/site-packages/dpctl/tensor/_manipulation_functions.py:379, in roll(X, shift, axis)
    375 res = dpt.empty(
    376     X.shape, dtype=X.dtype, usm_type=X.usm_type, sycl_queue=exec_q
    377 )
    378 dep_evs = _manager.submitted_events
--> 379 ht_e, roll_ev = ti._copy_usm_ndarray_for_roll_nd(
    380     src=X, dst=res, shifts=shifts, sycl_queue=exec_q, depends=dep_evs
    381 )
    382 _manager.add_event_pair(ht_e, roll_ev)
    383 return res

TypeError: _copy_usm_ndarray_for_roll_nd(): incompatible function arguments. The following argument types are supported:
    1. (src: dpctl::tensor::usm_ndarray, dst: dpctl::tensor::usm_ndarray, shifts: list[int], sycl_queue: dpctl.SyclQueue, depends: list[dpctl.SyclEvent] = []) -> tuple[dpctl.SyclEvent, dpctl.SyclEvent]

Invoked with: kwargs: src=usm_ndarray([0, 1, 2, 3, 4, 5, 6, 7, 8, 9]), dst=usm_ndarray([0, 1, 2, 3, 4, 5, 6, 7, 8, 9]), shifts=[1.0], sycl_queue=<dpctl.SyclQueue at 0x7f646250d8c0>, depends=[]

# works with signed integer:
dpt.roll(a, numpy.int64(1), axis=0)
# Out: usm_ndarray([9, 0, 1, 2, 3, 4, 5, 6, 7, 8])

A similar issue was present in NumPy, but resolved in 2.1.2 version.

[antonwolfy]

JFI, numpy resolved that by numpy#27431.

[ndgrigorian]

Strangely, I cannot replicate this on dev branch of dpctl.

In [18]: dpctl.__version__
Out[18]: '0.19.0dev0+124.g20064795dc'

In [19]: dpt.roll(a, np.uint64(1), axis=0)
Out[19]: usm_ndarray([9, 0, 1, 2, 3, 4, 5, 6, 7, 8])

[antonwolfy]

For me the same, I can't reproduce the initial problem.
But here is some corner case with big integer values which result into the same error for the latest dpctl:

dpctl.__version__
# Out: '0.19.0dev0+125.ga38af8cfd6'

numpy.roll(x, numpy.uint64(2**63+2))
# Out: array([2, 3, 0, 1])

numpy.roll(x, 2**100)
# Out: array([0, 1, 2, 3])

# but failed with dpctl:
x = dpt.arange(4)
dpt.roll(x, numpy.uint64(2**63+2))
---------------------------------------------------------------------------
TypeError                                 Traceback (most recent call last)
Cell In[18], line 1
----> 1 dpt.roll(x, numpy.uint64(2**63+2))

File /localdisk/work/antonvol/miniforge3/envs/dpctl_ext/lib/python3.12/site-packages/dpctl/tensor/_manipulation_functions.py:356, in roll(X, shift, axis)
    352 dep_evs = _manager.submitted_events
    353 res = dpt.empty(
    354     X.shape, dtype=X.dtype, usm_type=X.usm_type, sycl_queue=exec_q
    355 )
--> 356 hev, roll_ev = ti._copy_usm_ndarray_for_roll_1d(
    357     src=X,
    358     dst=res,
    359     shift=shift,
    360     sycl_queue=exec_q,
    361     depends=dep_evs,
    362 )
    363 _manager.add_event_pair(hev, roll_ev)
    364 return res

TypeError: _copy_usm_ndarray_for_roll_1d(): incompatible function arguments. The following argument types are supported:
    1. (src: dpctl::tensor::usm_ndarray, dst: dpctl::tensor::usm_ndarray, shift: int, sycl_queue: dpctl.SyclQueue, depends: list[dpctl.SyclEvent] = []) -> tuple[dpctl.SyclEvent, dpctl.SyclEvent]

Invoked with: kwargs: src=usm_ndarray([0, 1, 2, 3]), dst=usm_ndarray([0, 0, 0, 0]), shift=9223372036854775810, sycl_queue=<dpctl.SyclQueue at 0x7f524558e880>, depends=[]

# one more example:
dpt.roll(x, 2**100)
---------------------------------------------------------------------------
TypeError                                 Traceback (most recent call last)
Cell In[19], line 1
----> 1 dpt.roll(x, 2**100)

File /localdisk/work/antonvol/miniforge3/envs/dpctl_ext/lib/python3.12/site-packages/dpctl/tensor/_manipulation_functions.py:356, in roll(X, shift, axis)
    352 dep_evs = _manager.submitted_events
    353 res = dpt.empty(
    354     X.shape, dtype=X.dtype, usm_type=X.usm_type, sycl_queue=exec_q
    355 )
--> 356 hev, roll_ev = ti._copy_usm_ndarray_for_roll_1d(
    357     src=X,
    358     dst=res,
    359     shift=shift,
    360     sycl_queue=exec_q,
    361     depends=dep_evs,
    362 )
    363 _manager.add_event_pair(hev, roll_ev)
    364 return res

TypeError: _copy_usm_ndarray_for_roll_1d(): incompatible function arguments. The following argument types are supported:
    1. (src: dpctl::tensor::usm_ndarray, dst: dpctl::tensor::usm_ndarray, shift: int, sycl_queue: dpctl.SyclQueue, depends: list[dpctl.SyclEvent] = []) -> tuple[dpctl.SyclEvent, dpctl.SyclEvent]

Invoked with: kwargs: src=usm_ndarray([0, 1, 2, 3]), dst=usm_ndarray([0, 0, 0, 0]), shift=1267650600228229401496703205376, sycl_queue=<dpctl.SyclQueue at 0x7f524558e880>, depends=[]