Repository: numba/numba

Issue:
_lower_parfor_parallel function can mess up sequential_parfor_lowering flag disabling parallelization for the rest of the execution
<!--

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).

For more general "how do I do X?" type questions, please speak to us in real
time on https://gitter.im/numba/numba or post to the Numba mailing list
https://groups.google.com/a/continuum.io/forum/#!forum/numba-users.

-->

## Reporting a bug

<!--

Before submitting a bug report please ensure that you can check off these boxes:

-->

- [x] I am 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 below a minimal working reproducer (if you are unsure how
 to write one see http://matthewrocklin.com/blog/work/2018/02/28/minimal-bug-reports).

<!--

Please include details of the bug here, including, if applicable, what you
expected to happen!

-->

Repro:
```
class DummyType(types.Opaque):
    pass

dummy_type = DummyType("my_dummy")
register_model(DummyType)(OpaqueModel)

class Dummy(object):
    pass

@typeof_impl.register(Dummy)
def typeof_Dummy(val, c):
    return dummy_type

@unbox(DummyType)
def unbox_index(typ, obj, c):
    return NativeValue(c.context.get_dummy_value())

@overload_method(DummyType, "method1", jit_options={"parallel":True})
def _get_method1(obj, arr, func):
    def _foo(obj, arr, func):
        def baz(a, f):
            c = a.copy()
            c[numpy.isinf(a)] = numpy.nan
            return f(c)

        length = len(arr)
        output_arr = numpy.empty(length, dtype=numpy.float64)
        for i in prange(length):
            output_arr[i] = baz(arr[i], func)
        for i in prange(length - 1):
            output_arr[i] += baz(arr[i], func)
        return output_arr
    return _foo

@njit
def bar(v):
    return v.mean()

@njit
def test1(d):
    return d.method1(numpy.array([[1.0,2.0,3.0],[4.0,5.0,6.0]]), bar)

print(numba.parfor.sequential_parfor_lowering)
test1(Dummy())
print(numba.parfor.sequential_parfor_lowering)
```
Prints
```
False
...some parfor warnings...
True
```

There are actually multiple issues:
* If `_create_gufunc_for_parfor_body` https://github.com/numba/numba/blob/master/numba/npyufunc/parfor.py#L239 throws an error `sequential_parfor_lowering` flag wont be restored
* The error itself looks like this
```
Traceback (most recent call last):
  File "/home/ibutygin/src/numba/numba/numba/errors.py", line 717, in new_error_context
    yield
  File "/home/ibutygin/src/numba/numba/numba/lowering.py", line 288, in lower_block
    self.lower_inst(inst)
  File "/home/ibutygin/src/numba/numba/numba/lowering.py", line 476, in lower_inst
    func(self, inst)
  File "/home/ibutygin/src/numba/numba/numba/npyufunc/parfor.py", line 241, in _lower_parfor_parallel
    bool(alias_map), index_var_typ, parfor.races)
  File "/home/ibutygin/src/numba/numba/numba/npyufunc/parfor.py", line 833, in _create_gufunc_for_parfor_body
    replace_var_with_array(races, loop_body, typemap, lowerer.fndesc.calltypes)
  File "/home/ibutygin/src/numba/numba/numba/npyufunc/parfor.py", line 1206, in replace_var_with_array
    replace_var_with_array_internal(vars, loop_body, typemap, calltypes)
  File "/home/ibutygin/src/numba/numba/numba/npyufunc/parfor.py", line 1203, in replace_var_with_array_internal
    block.body = replace_var_with_array_in_block(vars, block, typemap, calltypes)
  File "/home/ibutygin/src/numba/numba/numba/npyufunc/parfor.py", line 1189, in replace_var_with_array_in_block
    setitem_node = ir.SetItem(inst.target, const_var, inst.value, inst.loc)
  File "/home/ibutygin/src/numba/numba/numba/ir.py", line 573, in __init__
    assert isinstance(value, Var)
AssertionError
```
* Error get silently ignored if we use `@overload_method` (but it will fail compilation if we make `_foo` a free funtion)

Tested on `release0.47` branch and master (`40fb2d06859a782db9463ac6d7eae8405eb6a81d`)


Discussion:
I can make a quick fix, wrapping `_create_gufunc_for_parfor_body` into `try..finally` but other issues are outside of my current competence
Whilst contrived, could this be used to reproduce and check the fix in #5114 ?
```python
from numba import njit, prange
import numba.parfors.parfor

from numba.core.compiler_machinery import register_pass, FunctionPass
from numba.core.compiler import CompilerBase, DefaultPassBuilder
from numba.core.typed_passes import IRLegalization

@register_pass(mutates_CFG=True, analysis_only=False)
class BreakParfors(FunctionPass):
    _name = "break_parfors"

    def __init__(self):
        FunctionPass.__init__(self)

    def run_pass(self, state):
        for blk in state.func_ir.blocks.values():
            for stmt in blk.body:
                if isinstance(stmt, numba.parfors.parfor.Parfor):
                    # races should be a set(), that list is iterable permits it to
                    # get through to the _create_gufunc_for_parfor_body routine
                    # at which point it needs to be a set so e.g. set.difference
                    # can be computed, this therefore creates an error in the
                    # right location
                    stmt.races = []
            return True


class BreakParforsCompiler(CompilerBase):

    def define_pipelines(self):
        pm = DefaultPassBuilder.define_nopython_pipeline(self.state)
        pm.add_pass_after(BreakParfors, IRLegalization)
        pm.finalize()
        return [pm]


@njit(parallel=True, pipeline_class=BreakParforsCompiler)
def foo():
    x = 1
    for _ in prange(1):
        x += 1
    return x


print(numba.parfors.parfor.sequential_parfor_lowering)
try:
    print(foo())
except Exception as e:
    print("Exception was raised:", e)
print(numba.parfors.parfor.sequential_parfor_lowering)
```

Resolving patch:
diff --git a/numba/parfors/parfor_lowering.py b/numba/parfors/parfor_lowering.py
--- a/numba/parfors/parfor_lowering.py
+++ b/numba/parfors/parfor_lowering.py
@@ -226,10 +226,12 @@ def _lower_parfor_parallel(lowerer, parfor):
     for l in parfor.loop_nests[1:]:
         assert typemap[l.index_variable.name] == index_var_typ
     numba.parfors.parfor.sequential_parfor_lowering = True
-    func, func_args, func_sig, redargstartdim, func_arg_types = _create_gufunc_for_parfor_body(
-        lowerer, parfor, typemap, typingctx, targetctx, flags, {},
-        bool(alias_map), index_var_typ, parfor.races)
-    numba.parfors.parfor.sequential_parfor_lowering = False
+    try:
+        func, func_args, func_sig, redargstartdim, func_arg_types = _create_gufunc_for_parfor_body(
+            lowerer, parfor, typemap, typingctx, targetctx, flags, {},
+            bool(alias_map), index_var_typ, parfor.races)
+    finally:
+        numba.parfors.parfor.sequential_parfor_lowering = False
 
     # get the shape signature
     func_args = ['sched'] + func_args


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
@@ -30,8 +30,11 @@
                             get_definition, is_getitem, is_setitem,
                             index_var_of_get_setitem)
 from numba.np.unsafe.ndarray import empty_inferred as unsafe_empty
-from numba.core.compiler import compile_isolated, Flags
 from numba.core.bytecode import ByteCodeIter
+from numba.core.compiler import (compile_isolated, Flags, CompilerBase,
+                                 DefaultPassBuilder)
+from numba.core.compiler_machinery import register_pass, FunctionPass
+from numba.core.typed_passes import IRLegalization
 from numba.tests.support import (TestCase, captured_stdout, MemoryLeakMixin,
                       override_env_config, linux_only, tag,
                       skip_parfors_unsupported, _32bit, needs_blas,
@@ -3098,6 +3101,61 @@ def test_impl():
 
         self.check(test_impl)
 
+    @skip_parfors_unsupported
+    def test_no_state_change_in_gufunc_lowering_on_error(self):
+        # tests #5098, if there's an exception arising in gufunc lowering the
+        # sequential_parfor_lowering global variable should remain as False on
+        # stack unwind.
+
+        @register_pass(mutates_CFG=True, analysis_only=False)
+        class BreakParfors(FunctionPass):
+            _name = "break_parfors"
+
+            def __init__(self):
+                FunctionPass.__init__(self)
+
+            def run_pass(self, state):
+                for blk in state.func_ir.blocks.values():
+                    for stmt in blk.body:
+                        if isinstance(stmt, numba.parfors.parfor.Parfor):
+                            # races should be a set(), that list is iterable
+                            # permits it to get through to the
+                            # _create_gufunc_for_parfor_body routine at which
+                            # point it needs to be a set so e.g. set.difference
+                            # can be computed, this therefore creates an error
+                            # in the right location.
+                            stmt.races = []
+                    return True
+
+
+        class BreakParforsCompiler(CompilerBase):
+
+            def define_pipelines(self):
+                pm = DefaultPassBuilder.define_nopython_pipeline(self.state)
+                pm.add_pass_after(BreakParfors, IRLegalization)
+                pm.finalize()
+                return [pm]
+
+
+        @njit(parallel=True, pipeline_class=BreakParforsCompiler)
+        def foo():
+            x = 1
+            for _ in prange(1):
+                x += 1
+            return x
+
+        # assert default state for global
+        self.assertFalse(numba.parfors.parfor.sequential_parfor_lowering)
+
+        with self.assertRaises(errors.LoweringError) as raises:
+            foo()
+
+        self.assertIn("'list' object has no attribute 'difference'",
+                      str(raises.exception))
+
+        # assert state has not changed
+        self.assertFalse(numba.parfors.parfor.sequential_parfor_lowering)
+
     @skip_parfors_unsupported
     def test_issue_5098(self):
         class DummyType(types.Opaque):
@@ -3153,7 +3211,6 @@ def test1(d):
             numba.parfors.parfor.sequential_parfor_lowering = save_state
 
 
-
 @skip_parfors_unsupported
 class TestParforsDiagnostics(TestParforsBase):