Repository: numba/numba
Pull request: Allow arbitrary walk-back in reduction nodes to find inplace_binop.
Resolves #7344 

Previously, there could only be one extra assignment after the inplace_binop node for a reduction before the assignment to the reduction variable.  In this issue, there was an additional assignment.  The solution is to start at the enf with the reduction variable and walk backwards through the reduce nodes so long as they are all assignments and link with each other until we get to inplace_binop.

Problem statement:
incoherent behaviour with prange
First, thanks a lot for this very useful job !

Here is the problem: I slightly change the example of prange given in the documentation and get a very strange behaviour. Here is my code
````python
from numba import njit, prange
@njit(parallel=True)
def prange_test(A, cond):
    s = 0
    for i in prange(A.shape[0]):
        if cond[i]:
            s += 1
    return s
````
Here are my inputs
````python
A = np.random.randint(0,10,100)
cond = np.random.randint(0,2,len(A)).astype('bool')
````
In this case scenario changing njit option from True to False do not affect the result of `prange_test(A, cond)` (which is what I expect)

However, setting `s=1` as a starting condition in `prange_test()` definition make `prange_test(A, cond)` results to be different depending of the parallel value !
If ``parallel==False``, the result is ok;  If ``parallel==True``, the results is larger of the quantity: 1 x number_threads +1.
Setting `s=2` as a starting condition in `prange_test()` definition makes the result to be larger of quantity 2 x number_of_threads +2.

This strange behavior disappears if there is no check of the ``cond`` value in `prange_test(A, cond)` (it returns the length of A+initial s value) whatever the initial value of ``s`` 
It seems inconsistent as it gives wrong results and threads dependent. I do not understand why checking  ``cond' creates inconsistencies.
Thanks for the help !

Hints from issue discussion:
I can reproduce the issue, and I think it's a bug - as far as I can tell from the docs (https://numba.readthedocs.io/en/latest/user/parallel.html#explicit-parallel-loops), this should be a supported case:

>  reduction is inferred automatically if a variable is updated by a binary function/operator using its previous value in the loop body. The initial value of the reduction is inferred automatically for the `+=`, `-=`, `*=`, and `/=` operators. 

A copy/pastable reproducer and test:

```python
from numba import njit, prange
import numpy as np


@njit(parallel=True)
def prange_init_0(A, cond):
    s = 0
    for i in prange(A.shape[0]):
        if cond[i]:
            s += 1
    return s


@njit(parallel=True)
def prange_init_1(A, cond):
    s = 1
    for i in prange(A.shape[0]):
        if cond[i]:
            s += 1
    return s


np.random.seed(1)
A = np.random.randint(0, 10, 100)
cond = np.random.randint(0, 2, len(A)).astype('bool')

init_0 = prange_init_0(A, cond)
init_1 = prange_init_1(A, cond)

print(init_0)
print(init_1)

print("Checking sanity...")
np.testing.assert_equal(init_0 + 1, init_1)
print("... Sanity checks passed!")
```

presently produces with master:

```
$ python repro.py
OMP: Info #276: omp_set_nested routine deprecated, please use omp_set_max_active_levels instead.
49
62
Checking sanity...
Traceback (most recent call last):
  File "/home/gmarkall/numbadev/issues/7344/repro.py", line 34, in <module>
    np.testing.assert_equal(init_0 + 1, init_1)
  File "/home/gmarkall/miniconda3/envs/numbanp120/lib/python3.9/site-packages/numpy/testing/_private/utils.py", line 423, in assert_equal
    raise AssertionError(msg)
AssertionError: 
Items are not equal:
 ACTUAL: 50
 DESIRED: 62
```

on my 12-core system, mirroring the behaviour reported here.

@DrTodd13 Could you have a look at this please?
The `if` condition seems to maybe be confusing the analysis for the reduction - If I change `prange_init_1` to:

```python
@njit(parallel=True)
def prange_init_1(A, cond):
    s = 1
    for i in prange(A.shape[0]):
        val = 0
        if cond[i]:
            val += 1
        s += val
    return s
```

then the correct result is obtained. @lanougue does this pattern provide a workaround for your code for the time being?
Hi @lanougue , could I try working on this? 
> Hi @lanougue , could I try working on this?

Anyone is free to contribute to Numba and work on any part of it, so you are welcome to to work on things that interest you!

However, I would note that this will be a difficult issue for a first-time contributor, as it will need a strong understanding of a lot of Numba internals - I would suggest that if you're keen to contribute to Numba, looking through the list of [good first issues](https://github.com/numba/numba/labels/good%20first%20issue) and picking one from there that seems interesting will be a more manageable task to begin with.
> The `if` condition seems to maybe be confusing the analysis for the reduction - If I change `prange_init_1` to:
> 
> ```python
> @njit(parallel=True)
> def prange_init_1(A, cond):
>     s = 1
>     for i in prange(A.shape[0]):
>         val = 0
>         if cond[i]:
>             val += 1
>         s += val
>     return s
> ```
> 
> then the correct result is obtained. @lanougue does this pattern provide a workaround for your code for the time being?

Thanks for the workaround, I will give a look. Thanks for the reactivity !

Whoever wants to look at this numba bug is welcome. My python level is clearly not good enough for that ! ;)
@lanougue For a quick fix, try the following.  I know it shouldn't matter if you do s = 1 at the beginning or s+= 1 at the end but given the bug it does matter.  We can do something here for += and *= but for more esoteric operators we may not know the identity value and so have to use the initial value as the identity value.  The bug is that in this case it doesn't recognize the operator as addition and so doesn't know the identity is 0 and so using the previous value of the variable.  The complex reason is that there's a hidden array of reduction values (each worker gets its own) and that array is currently initialized with the value of the reduction variable when entering the loop but that needs to be the identity value.

```
from numba import njit, prange
@njit(parallel=True)
def prange_test(A, cond):
    s = 0
    for i in prange(A.shape[0]):
        if cond[i]:
            s += 1
    s += 1
    return s
```

Commit message: Fix typo.

Patch to review:
diff --git a/numba/parfors/parfor.py b/numba/parfors/parfor.py
index afb7db8a183..cfc0afa627c 100644
--- a/numba/parfors/parfor.py
+++ b/numba/parfors/parfor.py
@@ -3566,15 +3566,17 @@ def get_reduction_init(nodes):
     Get initial value for known reductions.
     Currently, only += and *= are supported.
     """
-    require(len(nodes) >=1)
-    # there could be an extra assignment after the reduce node
+    require(len(nodes) >= 1)
+    # there could be multiple extra assignments after the reduce node
     # See: test_reduction_var_reuse
-    if isinstance(nodes[-1].value, ir.Var):
-        require(len(nodes) >=2)
-        require(nodes[-2].target.name == nodes[-1].value.name)
-        acc_expr = nodes[-2].value
-    else:
-        acc_expr = nodes[-1].value
+    cur_index = 1
+    acc_expr = nodes[-cur_index].value
+    while isinstance(acc_expr, ir.Var):
+        require(len(nodes) >= cur_index + 1)
+        require(nodes[-(cur_index + 1)].target.name == nodes[-cur_index].value.name)
+        cur_index += 1
+        acc_expr = nodes[-cur_index].value
+
     require(isinstance(acc_expr, ir.Expr) and acc_expr.op=='inplace_binop')
     if acc_expr.fn == operator.iadd or acc_expr.fn == operator.isub:
         return 0, acc_expr.fn
diff --git a/numba/tests/test_parfors.py b/numba/tests/test_parfors.py
index 4a93033e81d..d84278e72f4 100644
--- a/numba/tests/test_parfors.py
+++ b/numba/tests/test_parfors.py
@@ -1564,6 +1564,16 @@ def test_impl(n):
             return acc
         self.check(test_impl, 16)
 
+    def test_non_identity_initial(self):
+        # issue #7344
+        def test_impl(A, cond):
+            s = 1
+            for i in prange(A.shape[0]):
+                if cond[i]:
+                    s += 1
+            return s
+        self.check(test_impl, np.ones(10), np.ones(10).astype('bool'))
+
     def test_two_d_array_reduction_reuse(self):
         def test_impl(n):
             shp = (13, 17)


Code review comments:
diff_hunk: @@ -3566,15 +3566,17 @@ def get_reduction_init(nodes):
     Get initial value for known reductions.
     Currently, only += and *= are supported.
     """
-    require(len(nodes) >=1)
-    # there could be an extra assignment after the reduce node
+    require(len(nodes) >= 1)
+    # there could be multiple extra assignments after the reduce node
     # See: test_reduction_var_reuse
-    if isinstance(nodes[-1].value, ir.Var):
-        require(len(nodes) >=2)
-        require(nodes[-2].target.name == nodes[-1].value.name)
-        acc_expr = nodes[-2].value
-    else:
-        acc_expr = nodes[-1].value
+    cur_index = 1
+    acc_expr = nodes[-cur_index].value
+    while isinstance(acc_expr, ir.Var):
+        require(len(nodes) >= cur_index + 1)
+        require(nodes[-(cur_index + 1)].target.name == nodes[-cur_index].value.name)
+        cur_index += 1
+        acc_expr = nodes[-cur_index].value
path: numba/parfors/parfor.py
text: @user2:
It seems like this might be relying on a naming pattern? The following variant on the original MWR breaks for variable `t`:
```python
from numba import njit, prange
import numpy as np

@user1(parallel=True)
def test_impl(A, cond):
    s = 1
    t = 10
    for i in prange(A.shape[0]):
        if cond[i]:
            s += 1
            t += 1
        else:
            s += 2
    return s, t

print(test_impl(np.ones(10), (np.arange(10) % 2).astype('bool')))
```