Repository: numba/numba

Issue:
Issue with allocating ndarrays inside `prange`
I recently experienced some problems when allocating ndarrays inside `prange`. Here is a simple reproducer (it is quite a silly example, but this is because it has been reduced to the shortest possible example). Note that this uses numba version `0.36.0.dev+100.g4afaf2b`, but that the same issue appears for numba version `0.35`.

```python
from numba import njit, prange
import numpy as np

@njit(parallel=True)
def add_one_with_1darray(A,N):
    for iz in prange(N):
        x = np.ones( 1 )
        A[iz] += x[ 0 ]

@njit(parallel=True)
def add_one_with_2darray(A,N):
    for iz in prange(N):
        x = np.ones( (1, 1) )
        A[iz] += x[0,0]

A = np.zeros( (100,), dtype=np.complex128 )
A += np.random.rand(100)

add_one_with_1darray(A,100) # Works                                             
add_one_with_2darray(A,100) # Fails to compile 
```
The issue disappears if I remove `parallel=True` and `prange`.

Is this an issue with the way the code is written, or is this a bug with `numba`?

Thanks in advance for your help,
Issue with allocating ndarrays inside `prange`
I recently experienced some problems when allocating ndarrays inside `prange`. Here is a simple reproducer (it is quite a silly example, but this is because it has been reduced to the shortest possible example). Note that this uses numba version `0.36.0.dev+100.g4afaf2b`, but that the same issue appears for numba version `0.35`.

```python
from numba import njit, prange
import numpy as np

@njit(parallel=True)
def add_one_with_1darray(A,N):
    for iz in prange(N):
        x = np.ones( 1 )
        A[iz] += x[ 0 ]

@njit(parallel=True)
def add_one_with_2darray(A,N):
    for iz in prange(N):
        x = np.ones( (1, 1) )
        A[iz] += x[0,0]

A = np.zeros( (100,), dtype=np.complex128 )
A += np.random.rand(100)

add_one_with_1darray(A,100) # Works                                             
add_one_with_2darray(A,100) # Fails to compile 
```
The issue disappears if I remove `parallel=True` and `prange`.

Is this an issue with the way the code is written, or is this a bug with `numba`?

Thanks in advance for your help,


Discussion:
Thanks for the report, I can reproduce on the most-patched version of this feature. I also tested `0d` arrays, these work ok.

Backtrace is:
```
$ python issue2587.py 
Traceback (most recent call last):
  File "issue2587.py", line 27, in <module>
    add_one_with_2darray(A,100) # Fails to compile 
  File "<path>/numba/dispatcher.py", line 330, in _compile_for_args
    raise e
  File "<path>/numba/dispatcher.py", line 307, in _compile_for_args
    return self.compile(tuple(argtypes))
  File "<path>/numba/dispatcher.py", line 579, in compile
    cres = self._compiler.compile(args, return_type)
  File "<path>/numba/dispatcher.py", line 80, in compile
    flags=flags, locals=self.locals)
  File "<path>/numba/compiler.py", line 763, in compile_extra
    return pipeline.compile_extra(func)
  File "<path>/numba/compiler.py", line 360, in compile_extra
    return self._compile_bytecode()
  File "<path>/numba/compiler.py", line 722, in _compile_bytecode
    return self._compile_core()
  File "<path>/numba/compiler.py", line 709, in _compile_core
    res = pm.run(self.status)
  File "<path>/numba/compiler.py", line 246, in run
    raise patched_exception
  File "<path>/numba/compiler.py", line 238, in run
    stage()
  File "<path>/numba/compiler.py", line 644, in stage_nopython_backend
    self._backend(lowerfn, objectmode=False)
  File "<path>/numba/compiler.py", line 599, in _backend
    lowered = lowerfn()
  File "<path>/numba/compiler.py", line 586, in backend_nopython_mode
    self.flags)
  File "<path>/numba/compiler.py", line 882, in native_lowering_stage
    lower.lower()
  File "<path>/numba/lowering.py", line 135, in lower
    self.lower_normal_function(self.fndesc)
  File "<path>/numba/lowering.py", line 176, in lower_normal_function
    entry_block_tail = self.lower_function_body()
  File "<path>/numba/lowering.py", line 201, in lower_function_body
    self.lower_block(block)
  File "<path>/numba/lowering.py", line 216, in lower_block
    self.lower_inst(inst)
  File "<path>/numba/lowering.py", line 365, in lower_inst
    func(self, inst)
  File "<path>/numba/npyufunc/parfor.py", line 66, in _lower_parfor_parallel
    lowerer, parfor, typemap, typingctx, targetctx, flags, {})
  File "<path>/numba/npyufunc/parfor.py", line 605, in _create_gufunc_for_parfor_body
    locals)
  File "<path>/numba/compiler.py", line 777, in compile_ir
    lifted_from=lifted_from)
  File "<path>/numba/compiler.py", line 368, in compile_ir
    return self._compile_ir()
  File "<path>/numba/compiler.py", line 729, in _compile_ir
    return self._compile_core()
  File "<path>/numba/compiler.py", line 709, in _compile_core
    res = pm.run(self.status)
  File "<path>/numba/compiler.py", line 246, in run
    raise patched_exception
  File "<path>/numba/compiler.py", line 238, in run
    stage()
  File "<path>/numba/compiler.py", line 452, in stage_nopython_frontend
    self.locals)
  File "<path>/numba/compiler.py", line 865, in type_inference_stage
    infer.propagate()
  File "<path>/numba/typeinfer.py", line 844, in propagate
    raise errors[0]
  File "<path>/numba/typeinfer.py", line 137, in propagate
    constraint(typeinfer)
  File "<path>/numba/typeinfer.py", line 211, in __call__
    typeinfer.add_type(self.target, tup, loc=self.loc)
  File "<path>/numba/typeinfer.py", line 854, in add_type
    unified = tv.add_type(tp, loc=loc)
  File "<path>/numba/typeinfer.py", line 52, in add_type
    loc=loc)
numba.errors.TypingError: Failed at nopython (nopython mode backend)
Failed at nopython (nopython frontend)
No conversion from (int64 x 2) to int64 for '$parfor_index_tuple_var.104', defined at issue2587.py (19)
File "issue2587.py", line 19
[1] During: typing of tuple at issue2587.py (19)
[2] During: lowering "[LoopNest(index_variable = parfor_index.96, range = (0, N, 1))]{2: <ir.Block at issue2587.py (18)>}Var(parfor_index.96, issue2587.py (18))" at issue2587.py (18)
```
Seems like `A` isn't needed to trigger it.
```python
@njit(parallel=True)
def add_one_with_2darray(N):
    acc = 0
    for i in prange(N):
        x = np.ones((1, 1))
        acc += x[0, 0]
```
also fails the same way.

Ping @DrTodd13 @ehsantn @ninegua 
Thanks for the report, I can reproduce on the most-patched version of this feature. I also tested `0d` arrays, these work ok.

Backtrace is:
```
$ python issue2587.py 
Traceback (most recent call last):
  File "issue2587.py", line 27, in <module>
    add_one_with_2darray(A,100) # Fails to compile 
  File "<path>/numba/dispatcher.py", line 330, in _compile_for_args
    raise e
  File "<path>/numba/dispatcher.py", line 307, in _compile_for_args
    return self.compile(tuple(argtypes))
  File "<path>/numba/dispatcher.py", line 579, in compile
    cres = self._compiler.compile(args, return_type)
  File "<path>/numba/dispatcher.py", line 80, in compile
    flags=flags, locals=self.locals)
  File "<path>/numba/compiler.py", line 763, in compile_extra
    return pipeline.compile_extra(func)
  File "<path>/numba/compiler.py", line 360, in compile_extra
    return self._compile_bytecode()
  File "<path>/numba/compiler.py", line 722, in _compile_bytecode
    return self._compile_core()
  File "<path>/numba/compiler.py", line 709, in _compile_core
    res = pm.run(self.status)
  File "<path>/numba/compiler.py", line 246, in run
    raise patched_exception
  File "<path>/numba/compiler.py", line 238, in run
    stage()
  File "<path>/numba/compiler.py", line 644, in stage_nopython_backend
    self._backend(lowerfn, objectmode=False)
  File "<path>/numba/compiler.py", line 599, in _backend
    lowered = lowerfn()
  File "<path>/numba/compiler.py", line 586, in backend_nopython_mode
    self.flags)
  File "<path>/numba/compiler.py", line 882, in native_lowering_stage
    lower.lower()
  File "<path>/numba/lowering.py", line 135, in lower
    self.lower_normal_function(self.fndesc)
  File "<path>/numba/lowering.py", line 176, in lower_normal_function
    entry_block_tail = self.lower_function_body()
  File "<path>/numba/lowering.py", line 201, in lower_function_body
    self.lower_block(block)
  File "<path>/numba/lowering.py", line 216, in lower_block
    self.lower_inst(inst)
  File "<path>/numba/lowering.py", line 365, in lower_inst
    func(self, inst)
  File "<path>/numba/npyufunc/parfor.py", line 66, in _lower_parfor_parallel
    lowerer, parfor, typemap, typingctx, targetctx, flags, {})
  File "<path>/numba/npyufunc/parfor.py", line 605, in _create_gufunc_for_parfor_body
    locals)
  File "<path>/numba/compiler.py", line 777, in compile_ir
    lifted_from=lifted_from)
  File "<path>/numba/compiler.py", line 368, in compile_ir
    return self._compile_ir()
  File "<path>/numba/compiler.py", line 729, in _compile_ir
    return self._compile_core()
  File "<path>/numba/compiler.py", line 709, in _compile_core
    res = pm.run(self.status)
  File "<path>/numba/compiler.py", line 246, in run
    raise patched_exception
  File "<path>/numba/compiler.py", line 238, in run
    stage()
  File "<path>/numba/compiler.py", line 452, in stage_nopython_frontend
    self.locals)
  File "<path>/numba/compiler.py", line 865, in type_inference_stage
    infer.propagate()
  File "<path>/numba/typeinfer.py", line 844, in propagate
    raise errors[0]
  File "<path>/numba/typeinfer.py", line 137, in propagate
    constraint(typeinfer)
  File "<path>/numba/typeinfer.py", line 211, in __call__
    typeinfer.add_type(self.target, tup, loc=self.loc)
  File "<path>/numba/typeinfer.py", line 854, in add_type
    unified = tv.add_type(tp, loc=loc)
  File "<path>/numba/typeinfer.py", line 52, in add_type
    loc=loc)
numba.errors.TypingError: Failed at nopython (nopython mode backend)
Failed at nopython (nopython frontend)
No conversion from (int64 x 2) to int64 for '$parfor_index_tuple_var.104', defined at issue2587.py (19)
File "issue2587.py", line 19
[1] During: typing of tuple at issue2587.py (19)
[2] During: lowering "[LoopNest(index_variable = parfor_index.96, range = (0, N, 1))]{2: <ir.Block at issue2587.py (18)>}Var(parfor_index.96, issue2587.py (18))" at issue2587.py (18)
```
Seems like `A` isn't needed to trigger it.
```python
@njit(parallel=True)
def add_one_with_2darray(N):
    acc = 0
    for i in prange(N):
        x = np.ones((1, 1))
        acc += x[0, 0]
```
also fails the same way.

Ping @DrTodd13 @ehsantn @ninegua 

Resolving patch:
diff --git a/numba/parfor.py b/numba/parfor.py
--- a/numba/parfor.py
+++ b/numba/parfor.py
@@ -2138,8 +2138,8 @@ def parfor_typeinfer(parfor, typeinferer):
     save_blocks = typeinferer.blocks
     blocks = wrap_parfor_blocks(parfor)
     index_vars = [l.index_variable for l in parfor.loop_nests]
-    if len(parfor.loop_nests) > 1:
-        index_vars.append(parfor.index_var)
+    # no need to handle parfor.index_var (tuple of variables), since it will be
+    # assigned to a tuple from individual indices
     first_block = min(blocks.keys())
     loc = blocks[first_block].loc
     index_assigns = [ir.Assign(ir.Const(1, loc), v, loc) for v in index_vars]


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
@@ -889,6 +889,19 @@ def test_impl(A):
         # is updated before accumulation
         self.prange_tester(test_impl, np.random.ranf(4))
 
+    @skip_unsupported
+    def test_prange15(self):
+        # from issue 2587
+        # test parfor type inference when there is multi-dimensional indexing
+        def test_impl(N):
+            acc = 0
+            for i in range(N):
+                x = np.ones((1, 1))
+                acc += x[0, 0]
+            return acc
+
+        self.prange_tester(test_impl, 1024)
+
     @skip_unsupported
     def test_kde_example(self):
         def test_impl(X):