Repository: numba/numba

Issue:
ValueError: Argument types for wrap_index must match in @njit(parallel=True) when using int32 as index
This should work:

```py
import numba
import numpy as np

@numba.njit(parallel=True)
def test(data, indptr):
    sums = np.zeros((indptr.size - 1), dtype=data.dtype)
    for i in numba.prange(indptr.size - 1):
        start, end = indptr[i], indptr[i + 1]
        sums[i] = np.sum(data[start:end])
    return sums.reshape((indptr.size - 1, 1))

test(np.array([.3, .4, .5, .6]), np.array([0, 2, 4], dtype=np.int32))
```

But I get:

```pytb
Traceback (most recent call last):
  File "/home/phil/.PyCharm2019.3/config/scratches/scratch_1.py", line 12, in <module>
    test(np.array([.3, .4, .5, .6]),np.array([0, 2, 4], dtype=np.int32))
  File "/usr/lib/python3.8/site-packages/numba/dispatcher.py", line 401, in _compile_for_args
    error_rewrite(e, 'typing')
  File "/usr/lib/python3.8/site-packages/numba/dispatcher.py", line 344, in error_rewrite
    reraise(type(e), e, None)
  File "/usr/lib/python3.8/site-packages/numba/six.py", line 668, in reraise
    raise value.with_traceback(tb)
numba.errors.TypingError: Failed in nopython mode pipeline (step: nopython mode backend)
Failed in nopython mode pipeline (step: nopython frontend)
Invalid use of Function(<intrinsic wrap_index>) with argument(s) of type(s): (int32, int64)
 * parameterized
In definition 0:
    ValueError: Argument types for wrap_index must match
    raised from /usr/lib/python3.8/site-packages/numba/array_analysis.py:72
In definition 1:
    ValueError: Argument types for wrap_index must match
    raised from /usr/lib/python3.8/site-packages/numba/array_analysis.py:72
This error is usually caused by passing an argument of a type that is unsupported by the named function.
[1] During: resolving callee type: Function(<intrinsic wrap_index>)
[2] During: typing of call at /home/phil/.PyCharm2019.3/config/scratches/scratch_1.py (9)


File "scratch_1.py", line 9:
def test(data, indptr):
    <source elided>
        start, end = indptr[i], indptr[i + 1]
        sums[i] = np.sum(data[start:end])
        ^

[1] During: lowering "id=2[LoopNest(index_variable = parfor_index.50, range = (0, $34binary_subtract.15, 1))]{42: <ir.Block at /home/phil/.PyCharm2019.3/config/scratches/scratch_1.py (7)>}Var(parfor_index.50, scratch_1.py:7)" at /home/phil/.PyCharm2019.3/config/scratches/scratch_1.py (7)
```

Numba 0.47.0


Discussion:
Thanks for the report, I can reproduce. Looks like this is a bug in array analysis used by the parallel CPU target.

Resolving patch:
diff --git a/numba/array_analysis.py b/numba/array_analysis.py
--- a/numba/array_analysis.py
+++ b/numba/array_analysis.py
@@ -68,27 +68,16 @@ def wrap_index(typingctx, idx, size):
     where idx > size due to the way indices are calculated
     during slice/range analysis.
     """
-    if idx != size:
+    unified_ty = typingctx.unify_types(idx, size)
+    if not unified_ty:
         raise ValueError("Argument types for wrap_index must match")
 
     def codegen(context, builder, sig, args):
-        """
-        assert(len(args) == 2)
-        idx = args[0]
-        size = args[1]
-        rem = builder.srem(idx, size)
-        zero = llvmlite.ir.Constant(idx.type, 0)
-        is_negative = builder.icmp_signed('<', rem, zero)
-        wrapped_rem = builder.add(rem, size)
-        is_oversize = builder.icmp_signed('>=', wrapped_rem, size)
-        mod = builder.select(is_negative, wrapped_rem,
-                builder.select(is_oversize, rem, wrapped_rem))
-        return mod
-        """
-        idx = args[0]
-        size = args[1]
+        ll_unified_ty = context.get_data_type(unified_ty)
+        idx = builder.sext(args[0], ll_unified_ty)
+        size = builder.sext(args[1], ll_unified_ty)
         neg_size = builder.neg(size)
-        zero = llvmlite.ir.Constant(idx.type, 0)
+        zero = llvmlite.ir.Constant(ll_unified_ty, 0)
         idx_negative = builder.icmp_signed('<', idx, zero)
         pos_oversize = builder.icmp_signed('>=', idx, size)
         neg_oversize = builder.icmp_signed('<=', idx, neg_size)
@@ -97,7 +86,7 @@ def codegen(context, builder, sig, args):
         mod = builder.select(idx_negative, neg_res, pos_res)
         return mod
 
-    return signature(idx, idx, size), codegen
+    return signature(unified_ty, idx, size), codegen
 
 def wrap_index_literal(idx, size):
     if idx < 0:


Test patch:
diff --git a/numba/tests/test_array_analysis.py b/numba/tests/test_array_analysis.py
--- a/numba/tests/test_array_analysis.py
+++ b/numba/tests/test_array_analysis.py
@@ -993,6 +993,21 @@ def test_impl2(A, a):
         np.testing.assert_array_equal(
             njit(test_impl2, parallel=True)(A, a), test_impl2(A, a))
 
+    @skip_unsupported
+    def test_slice_dtype_issue_5056(self):
+        # see issue 5056
+
+        @njit(parallel=True)
+        def test_impl(data):
+            N = data.shape[0]
+            sums = np.zeros(N)
+            for i in prange(N):
+                sums[i] = np.sum(data[np.int32(0):np.int32(1)])
+            return sums
+
+        data = np.arange(10.)
+        np.testing.assert_array_equal(test_impl(data), test_impl.py_func(data))
+
 
 if __name__ == '__main__':
     unittest.main()