Repository: numba/numba

Issue:
Fail to compile `prange` with slice in variable resulting in attempt to `getitem` with wrong number of indices
Whilst attempting to minimize the test case given in https://numba.discourse.group/t/how-do-i-parallelize-this-code/1085/6?u=gmarkall I came across this issue. The following:

```python
import numpy as np
from numba import prange, njit


@njit(parallel=True)
def myFoo(x, par, numthreads):
    n_par = par.shape[0]
    n_x = len(x)
    result = np.zeros((n_par, n_x), dtype=np.float64)
    print(result.shape)
    chunklen = (len(x) + numthreads - 1) // numthreads

    for i in prange(numthreads):
        start = i * chunklen
        stop = (i + 1) * chunklen
        result[:, start:stop] = x[start:stop] * par[:]

    return result


x = np.array(np.arange(0, 12, 1.0))
par = np.array([1.0, 2.0, 3.0])
print(myFoo.py_func(x, par, 4))
print(myFoo(x, par, 4))
```

when run with `parallel=False` produces:

```
(3, 12)
[[ 0.  2.  6.  3.  8. 15.  6. 14. 24.  9. 20. 33.]
 [ 0.  2.  6.  3.  8. 15.  6. 14. 24.  9. 20. 33.]
 [ 0.  2.  6.  3.  8. 15.  6. 14. 24.  9. 20. 33.]]
(3, 12)
[[ 0.  2.  6.  3.  8. 15.  6. 14. 24.  9. 20. 33.]
 [ 0.  2.  6.  3.  8. 15.  6. 14. 24.  9. 20. 33.]
 [ 0.  2.  6.  3.  8. 15.  6. 14. 24.  9. 20. 33.]]
```

When run with `parallel=True` gives:

```
(3, 12)
[[ 0.  2.  6.  3.  8. 15.  6. 14. 24.  9. 20. 33.]
 [ 0.  2.  6.  3.  8. 15.  6. 14. 24.  9. 20. 33.]
 [ 0.  2.  6.  3.  8. 15.  6. 14. 24.  9. 20. 33.]]
OMP: Info #276: omp_set_nested routine deprecated, please use omp_set_max_active_levels instead.
Traceback (most recent call last):
  File "/home/gmarkall/numbadev/issues/discourse-1085/repro.py", line 24, in <module>
    print(myFoo(x, par, 4))
  File "/home/gmarkall/numbadev/numba/numba/core/dispatcher.py", line 468, in _compile_for_args
    error_rewrite(e, 'typing')
  File "/home/gmarkall/numbadev/numba/numba/core/dispatcher.py", line 409, in error_rewrite
    raise e.with_traceback(None)
numba.core.errors.TypingError: Failed in nopython mode pipeline (step: native lowering)
Failed in nopython mode pipeline (step: nopython frontend)
No implementation of function Function(<built-in function getitem>) found for signature:
 
 >>> getitem(array(float64, 1d, C), UniTuple(int64 x 2))
 
There are 22 candidate implementations:
   - Of which 20 did not match due to:
   Overload of function 'getitem': File: <numerous>: Line N/A.
     With argument(s): '(array(float64, 1d, C), UniTuple(int64 x 2))':
    No match.
   - Of which 2 did not match due to:
   Overload in function 'GetItemBuffer.generic': File: numba/core/typing/arraydecl.py: Line 166.
     With argument(s): '(array(float64, 1d, C), UniTuple(int64 x 2))':
    Rejected as the implementation raised a specific error:
      NumbaTypeError: cannot index array(float64, 1d, C) with 2 indices: UniTuple(int64 x 2)
  raised from /home/gmarkall/numbadev/numba/numba/core/typing/arraydecl.py:88

During: typing of intrinsic-call at /home/gmarkall/numbadev/issues/discourse-1085/repro.py (16)

File "repro.py", line 16:
def myFoo(x, par, numthreads):
    <source elided>
        stop = (i + 1) * chunklen
        result[:, start:stop] = x[start:stop] * par[:]
        ^

During: lowering "id=3[LoopNest(index_variable = parfor_index.50, range = (0, numthreads, 1))]{80: <ir.Block at /home/gmarkall/numbadev/issues/discourse-1085/repro.py (13)>}Var(parfor_index.50, repro.py:13)" at /home/gmarkall/numbadev/issues/discourse-1085/repro.py (13)
```

Could you have a look at this please @DrTodd13 ?


Discussion:

Resolving patch:
diff --git a/numba/parfors/parfor.py b/numba/parfors/parfor.py
--- a/numba/parfors/parfor.py
+++ b/numba/parfors/parfor.py
@@ -1763,7 +1763,27 @@ def run(self, blocks):
                         else:
                             # Handle A[:] = x
                             shape = equiv_set.get_shape(instr)
-                            if shape is not None:
+                            # Don't converted broadcasted setitems into parfors.
+                            if isinstance(index_typ, types.BaseTuple):
+                                # The sliced dims are those in the index that
+                                # are made of slices.  Count the numbers of slices
+                                # in the index tuple.
+                                sliced_dims = len(list(filter(
+                                    lambda x: isinstance(x, types.misc.SliceType),
+                                    index_typ.types)))
+                            elif isinstance(index_typ, types.misc.SliceType):
+                                # For singular indices there can be a bare slice
+                                # and if so there is one dimension being set.
+                                sliced_dims = 1
+                            else:
+                                sliced_dims = 0
+
+                            # Only create a parfor for this setitem if we know the
+                            # shape of the output and number of dimensions set is
+                            # equal to the number of dimensions on the right side.
+                            if (shape is not None and
+                                (not isinstance(value_typ, types.npytypes.Array) or
+                                 sliced_dims == value_typ.ndim)):
                                 new_instr = self._setitem_to_parfor(equiv_set,
                                         loc, target, index, value, shape=shape)
                                 self.rewritten.append(


Test patch:
diff --git a/numba/tests/test_parfors.py b/numba/tests/test_parfors.py
--- a/numba/tests/test_parfors.py
+++ b/numba/tests/test_parfors.py
@@ -3569,6 +3569,26 @@ def test_impl(flag):
         self.prange_tester(test_impl, True)
         self.prange_tester(test_impl, False)
 
+    def test_prange30(self):
+        # issue7675: broadcast setitem
+        def test_impl(x, par, numthreads):
+            n_par = par.shape[0]
+            n_x = len(x)
+            result = np.zeros((n_par, n_x), dtype=np.float64)
+            chunklen = (len(x) + numthreads - 1) // numthreads
+
+            for i in range(numthreads):
+                start = i * chunklen
+                stop = (i + 1) * chunklen
+                result[:, start:stop] = x[start:stop] * par[:]
+
+            return result
+
+        x = np.array(np.arange(0, 6, 1.0))
+        par = np.array([1.0, 2.0, 3.0])
+
+        self.prange_tester(test_impl, x, par, 2)
+
 
 @skip_parfors_unsupported
 class TestPrangeSpecific(TestPrangeBase):