Repository: numba/numba

Issue:
conditional is not evaluated correctly if parallel=True
<!--

Thanks for opening an issue! To help the Numba team handle your information
efficiently, please first ensure that there is no other issue present that
already describes the issue you have
(search at https://github.com/numba/numba/issues?&q=is%3Aissue).

-->

## Reporting a bug

<!--

Before submitting a bug report please ensure that you can check off these boxes:

-->

- [X] I have tried using the latest released version of Numba (most recent is
 visible in the change log (https://github.com/numba/numba/blob/master/CHANGE_LOG).
- [X] I have included a self contained code sample to reproduce the problem.
  i.e. it's possible to run as 'python bug.py'.
The self contained code sample is supplied both inline and as `bug.py` within the attached `bug.zip`.
[bug.zip](https://github.com/numba/numba/files/6512636/bug.zip)

<!--

Please include details of the bug here, including, if applicable, what you
expected to happen!

-->

The `else` path in a top level `if-else` in a `@njit(parallel=True)` compiled function with separate `prange`d for loops produces wrong results:
```python
import numpy as np
from numba import njit, prange
@njit(parallel=True)
def case_parallel(size,case):
    result = np.zeros((size,))
    if case == 1:
        for i in prange(size):
            result[i] += 1
    else:
        for i in prange(size):
            result[i] += 2
    return result.mean()
@njit(parallel=False)
def case_serial(size,case):
    result = np.zeros((size,))
    if case == 1:
        for i in prange(size):
            result[i] += 1
    else:
        for i in prange(size):
            result[i] += 2
    return result.mean()
for size in [1,10,100]:
    print(f'serial  (size={size:3}), should be 1.0: {case_serial(size=size,case=1)}')
    print(f'parallel(size={size:3}), should be 1.0: {case_parallel(size=size,case=1)}')
``` 
prints
```
serial  (size=  1), should be 1.0: 1.0
parallel(size=  1), should be 1.0: 3.0
serial  (size= 10), should be 1.0: 1.0
parallel(size= 10), should be 1.0: 3.0
serial  (size=100), should be 1.0: 1.0
parallel(size=100), should be 1.0: 3.0
```
As can be seen, this bug is not dependent on actual parallel computations.

I'm running numba 0.53.1 and numpy 1.19.4 on python 3.8.8 on Linux-3.10.0-1160.15.2.el7.x86_64-x86_64-with-glibc2.10



Discussion:
I am able to reproduce this issue when the threading layer is tbb. But, when the threading layer is omp, I see the code giving correct outputs. I was able to verify this on binder. You can verify it [here](https://mybinder.org/v2/gh/rawwar/NumbaTriaging/HEAD)
@rawwar thank you for looking into this. Looking at the discovery, this points to this maybe being a tbb bug.
I can reproduce with OpenMP and workqueue too, MWR:
```python
import numpy as np
from numba import njit, prange, threading_layer

@njit(parallel=True)
def case_parallel(size, case):
    result = np.zeros((size,))
    if case == 1:
        for i in prange(size):
            result[i] += 1
    else:
        for i in prange(size):
            result[i] += 2
    return result[0]

print(f'{case_parallel(3, 1)}, with threading_layer: {threading_layer()}')
print(f'{case_parallel.py_func(3, 1)}')
```
gives:
```
$ NUMBA_THREADING_LAYER=omp python issue7051.py 
3.0, with threading_layer: omp
1.0

$ NUMBA_THREADING_LAYER=workqueue python issue7051.py 
3.0, with threading_layer: workqueue
1.0
```

I think this is an issue with the parfor pass, specifically `prange` induced loop conversion, it seems like `parfor 0`, the second `prange` driven loop, is unconditionally executed. Looking at the Numba IR CFG before and after execution of `numba.parfors.parfor.ConvertLoopPass` confirms locally.

CC @DrTodd13 
I cannot reproduce on my end with any of the threading backends.
This should help debug (needs `conda install python-graphviz` to render the `dot` files).
```python
import numpy as np
from numba import njit, prange, threading_layer
from numba.core.compiler import CompilerBase, DefaultPassBuilder
from numba.core.compiler_machinery import FunctionPass, register_pass
from numba.core.typed_passes import PreParforPass, ParforPass

def gen_dump(name):
    @register_pass(mutates_CFG=False, analysis_only=True)
    class DumpNumbaCFG(FunctionPass):
        _name = "dump_numba_cfg_%s" % name

        def __init__(self):
            FunctionPass.__init__(self)

        def run_pass(self, state):
            state.func_ir.render_dot(filename_prefix=name).view()
            return False
    return DumpNumbaCFG


class DebugCompiler(CompilerBase):

    def define_pipelines(self):
        pm = DefaultPassBuilder.define_nopython_pipeline(self.state)
        pm.add_pass_after(gen_dump('before_parfors'), PreParforPass)
        pm.add_pass_after(gen_dump('after_parfors'), ParforPass)
        [print(x) for x in pm.passes]
        pm.finalize()
        return [pm]


@njit(parallel=True, pipeline_class=DebugCompiler)
def case_parallel(size, case):
    result = np.zeros((size,))
    if case == 1:
        for i in prange(size):
            result[i] += 1
    else:
        for i in prange(size):
            result[i] += 2
    return result[0]

print(f'{case_parallel(3, 1)}, with threading_layer: {threading_layer()}')
```
it renders the Numba IR CFG after both the `PreParforPass` and the `ParforPass`. Am testing issue with Python 3.8.5.
The reproducibility seems to be dependent on the python version. With the same version of numba 0.53.1 across all three threading layers (omp, tbb, workqueue), I get the bug using python 3.8.8, but not in 3.6.10.

Resolving patch:
diff --git a/numba/parfors/parfor.py b/numba/parfors/parfor.py
--- a/numba/parfors/parfor.py
+++ b/numba/parfors/parfor.py
@@ -2301,6 +2301,11 @@ def run(self, blocks):
                         and isinstance(inst.value, ir.Expr)
                         and inst.value.op == 'call'
                         and self._is_parallel_loop(inst.value.func.name, call_table)):
+                    # Here we've found a parallel loop, either prange or pndindex.
+                    # We create a parfor from this loop and then overwrite the contents
+                    # of the original loop header block to contain this parfor and then
+                    # a jump to the original loop exit block.  Other blocks in the
+                    # original loop are discarded.
                     body_labels = [ l for l in loop.body if
                                     l in blocks and l != loop.header ]
                     args = inst.value.args
@@ -2539,19 +2544,17 @@ def find_mask_from_size(size_var):
                                     equiv_set,
                                     ("prange", loop_kind, loop_replacing),
                                     pass_states.flags, races=races)
-                    # add parfor to entry block's jump target
-                    jump = blocks[entry].body[-1]
-                    jump.target = list(loop.exits)[0]
-                    blocks[jump.target].body.insert(0, parfor)
+
+                    blocks[loop.header].body = [parfor, ir.Jump(list(loop.exits)[0], loc)]
                     self.rewritten.append(dict(
                         old_loop=loop,
                         new=parfor,
                         reason='loop',
                     ))
                     # remove loop blocks from top level dict
-                    blocks.pop(loop.header)
                     for l in body_labels:
-                        blocks.pop(l)
+                        if l != loop.header:
+                            blocks.pop(l)
                     if config.DEBUG_ARRAY_OPT >= 1:
                         print("parfor from loop")
                         parfor.dump()


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
@@ -2835,6 +2835,20 @@ def comparer(a, b):
                            states,
                            check_arg_equality=[comparer])
 
+    @skip_parfors_unsupported
+    def test_issue7501(self):
+        def test_impl(size, case):
+            result = np.zeros((size,))
+            if case == 1:
+                for i in range(size):
+                    result[i] += 1
+            else:
+                for i in range(size):
+                    result[i] += 2
+            return result[0]
+
+        self.prange_tester(test_impl, 3, 1)
+
 
 @skip_parfors_unsupported
 @x86_only