Repository: numba/numba

Issue:
prange seems to exit for iterator before all threads finish. Do I missing some code like join()?
When use for iterator with prange, it seems sometimes the main thread will exit "for iterator" and exec next code before all threads finished.
An example as the best explanation:
[tt.py.txt](https://github.com/numba/numba/files/8115429/tt.py.txt)




Discussion:
Let me give a more detail explanation for my code:
We have a 4x4 chess board. Initial values are zero. Boarder squares not change.
For inner squares, we will do such operations: The new value of a square is the sum of 4 adjacent squares plus 1.
We do these operations for 2 stages, first we change value of all white squares, the for black.
The value of white square will only depend on black squares and vice versa. So each stage can be parallelized.

That is what my code did. But the result reveals that probably some operations on black square are executed before white squares are really finished.




My environment:
Win 10 64 bit
python 3.6.8  CPython
numba           0.53.1
numpy           1.19.5
Repaste the reproducer code here:

```python
import numpy as np
import numba
from math import *
from numba import threading_layer, set_num_threads, config
config.THREADING_LAYER = 'workqueue'

size = 4

@numba.jit(nogil=True, nopython = True)
def proc(u, i, even):
    for j in range((even + i + 1)%2 +1, size-1, 2):
        # print(i, j, even)
        u[i, j] = u[i, j+1] + u[i, j-1] + u[i+1,j ] + u[i-1, j] + 1




@numba.jit(nogil=True, parallel = True)
def demo():
    u = np.zeros((size,size))
    for i in numba.prange(1, size-1):
        proc(u, i, 0)
    for i in numba.prange(1, size-1):
        proc(u, i, 1)
    return u

expect = np.array([[0,0,0,0],[0, 1,3,0],[0,3,1,0],[0,0,0,0]])

for r in range(30):
    res = demo()
    if (res != expect).any():
        print(res)
    else:
        print('ok')
```
I can confirm that switching off parfors (e.g. `parallel=False`) would passes the check in the reproducer
@iynaur, this is happening because the `prange` loops are fused together by optimizations in Numba but it should not have because of the second loop depends on the first. As a workaround, you can disable loop fusion by:

```python
@numba.jit(nogil=True, parallel=dict(prange=True, fusion=False))
                                ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
```

@DrTodd13, This looks like a problem of invalid loop fusion. Parfors cannot see that the two loops have dependency. Any idea how parfors can detect this problem?


@sklam Thanks for the explanation and workaround.

Resolving patch:
diff --git a/numba/parfors/parfor.py b/numba/parfors/parfor.py
--- a/numba/parfors/parfor.py
+++ b/numba/parfors/parfor.py
@@ -530,6 +530,7 @@ def __init__(
             equiv_set,
             pattern,
             flags,
+            *,  #only specify the options below by keyword
             no_sequential_lowering=False,
             races=set()):
         super(Parfor, self).__init__(
@@ -2860,13 +2861,19 @@ def run(self):
             maximize_fusion(self.func_ir, self.func_ir.blocks, self.typemap,
                                                             up_direction=False)
             dprint_func_ir(self.func_ir, "after maximize fusion down")
-            self.fuse_parfors(self.array_analysis, self.func_ir.blocks)
+            self.fuse_parfors(self.array_analysis,
+                              self.func_ir.blocks,
+                              self.func_ir,
+                              self.typemap)
             dprint_func_ir(self.func_ir, "after first fuse")
             # push non-parfors up
             maximize_fusion(self.func_ir, self.func_ir.blocks, self.typemap)
             dprint_func_ir(self.func_ir, "after maximize fusion up")
             # try fuse again after maximize
-            self.fuse_parfors(self.array_analysis, self.func_ir.blocks)
+            self.fuse_parfors(self.array_analysis,
+                              self.func_ir.blocks,
+                              self.func_ir,
+                              self.typemap)
             dprint_func_ir(self.func_ir, "after fusion")
             # remove dead code after fusion to remove extra arrays and variables
             simplify(self.func_ir, self.typemap, self.calltypes, self.metadata["parfors"])
@@ -2958,7 +2965,7 @@ def _mk_parfor_loops(self, size_vars, scope, loc):
         """
         return _mk_parfor_loops(self.typemap, size_vars, scope, loc)
 
-    def fuse_parfors(self, array_analysis, blocks):
+    def fuse_parfors(self, array_analysis, blocks, func_ir, typemap):
         for label, block in blocks.items():
             equiv_set = array_analysis.get_equiv_set(label)
             fusion_happened = True
@@ -2976,32 +2983,32 @@ def fuse_parfors(self, array_analysis, blocks):
                         stmt.equiv_set = equiv_set
                         next_stmt.equiv_set = equiv_set
                         fused_node, fuse_report = try_fuse(equiv_set, stmt, next_stmt,
-                            self.metadata["parfors"])
+                            self.metadata["parfors"], func_ir, typemap)
                         # accumulate fusion reports
                         self.diagnostics.fusion_reports.append(fuse_report)
                         if fused_node is not None:
                             fusion_happened = True
                             self.diagnostics.fusion_info[stmt.id].extend([next_stmt.id])
                             new_body.append(fused_node)
-                            self.fuse_recursive_parfor(fused_node, equiv_set)
+                            self.fuse_recursive_parfor(fused_node, equiv_set, func_ir, typemap)
                             i += 2
                             continue
                     new_body.append(stmt)
                     if isinstance(stmt, Parfor):
-                        self.fuse_recursive_parfor(stmt, equiv_set)
+                        self.fuse_recursive_parfor(stmt, equiv_set, func_ir, typemap)
                     i += 1
                 new_body.append(block.body[-1])
                 block.body = new_body
         return
 
-    def fuse_recursive_parfor(self, parfor, equiv_set):
+    def fuse_recursive_parfor(self, parfor, equiv_set, func_ir, typemap):
         blocks = wrap_parfor_blocks(parfor)
         maximize_fusion(self.func_ir, blocks, self.typemap)
         dprint_func_ir(self.func_ir, "after recursive maximize fusion down", blocks)
         arr_analysis = array_analysis.ArrayAnalysis(self.typingctx, self.func_ir,
                                                 self.typemap, self.calltypes)
         arr_analysis.run(blocks, equiv_set)
-        self.fuse_parfors(arr_analysis, blocks)
+        self.fuse_parfors(arr_analysis, blocks, func_ir, typemap)
         unwrap_parfor_blocks(parfor)
 
 
@@ -3931,7 +3938,7 @@ def get_parfor_writes(parfor):
 
 FusionReport = namedtuple('FusionReport', ['first', 'second', 'message'])
 
-def try_fuse(equiv_set, parfor1, parfor2, metadata):
+def try_fuse(equiv_set, parfor1, parfor2, metadata, func_ir, typemap):
     """try to fuse parfors and return a fused parfor, otherwise return None
     """
     dprint("try_fuse: trying to fuse \n", parfor1, "\n", parfor2)
@@ -3980,9 +3987,16 @@ def get_user_varname(v):
             report = FusionReport(parfor1.id, parfor2.id, msg % i)
             return None, report
 
+    func_ir._definitions = build_definitions(func_ir.blocks)
     # TODO: make sure parfor1's reduction output is not used in parfor2
     # only data parallel loops
-    if has_cross_iter_dep(parfor1) or has_cross_iter_dep(parfor2):
+    p1_cross_dep, p1_ip, p1_ia, p1_non_ia = has_cross_iter_dep(parfor1, func_ir, typemap)
+    if not p1_cross_dep:
+        p2_cross_dep = has_cross_iter_dep(parfor2, func_ir, typemap, p1_ip, p1_ia, p1_non_ia)[0]
+    else:
+        p2_cross_dep = True
+
+    if p1_cross_dep or p2_cross_dep:
         dprint("try_fuse: parfor cross iteration dependency found")
         msg = ("- fusion failed: cross iteration dependency found "
                 "between loops #%s and #%s")
@@ -4077,28 +4091,236 @@ def remove_duplicate_definitions(blocks, nameset):
     return
 
 
-def has_cross_iter_dep(parfor):
-    # we conservatively assume there is cross iteration dependency when
-    # the parfor index is used in any expression since the expression could
-    # be used for indexing arrays
+def has_cross_iter_dep(
+        parfor,
+        func_ir,
+        typemap,
+        index_positions=None,
+        indexed_arrays=None,
+        non_indexed_arrays=None):
+    # We should assume there is cross iteration dependency unless we can
+    # prove otherwise.  Return True if there is a cross-iter dependency
+    # that should prevent fusion, False if fusion is okay.
+    # Also returns index_positions, indexed_arrays, and non_indexed_arrays
+    # who purpose is described below so that subsequent additional
+    # has_cross_iter_dep calls for other parfors can build on the same
+    # data structures to make sure that the array accesses generate no
+    # cross-iter dependencies both within a parfor but also across parfors.
+
     # TODO: make it more accurate using ud-chains
-    indices = {l.index_variable for l in parfor.loop_nests}
+
+    # Get the index variable used by this parfor.
+    # This will hold other variables with equivalent value, e.g., a = index_var
+    indices = {l.index_variable.name for l in parfor.loop_nests}
+    # This set will store variables that are (potentially recursively)
+    # defined in relation to an index variable, e.g., a = index_var + 1.
+    # A getitem that uses an index variable from this set will be considered
+    # as potentially having a cross-iter dependency and so won't fuse.
+    derived_from_indices = set()
+    # For the first parfor considered for fusion, the latter 3 args will be None
+    # and initialized to empty.  For the second parfor, the structures from the
+    # previous parfor are passed in so that cross-parfor violations of the
+    # below comments can prevent fusion.
+    #
+    # index_positions keeps track of which index positions have had an index
+    # variable used for them and which ones haven't for each possible array
+    # dimensionality.  After the first array access is seen, if subsequent
+    # ones use a parfor index for a different dimension then we conservatively
+    # say that we can't fuse.  For example, if a 2D array is accessed with
+    # a[parfor_index, 0] then index_positions[2] will be (True, False) and
+    # if a[0, parfor_index] happens later which is (False, True) then this
+    # conflicts with the previous value and will prevent fusion.
+    #
+    # indexed_arrays records arrays that are accessed with at least one
+    # parfor index.  If such an array is later accessed with indices that
+    # don't include a parfor index then conservatively assume we can't fuse.
+    #
+    # non_indexed_arrays holds arrays that are indexed without a parfor index.
+    # If an array first accessed without a parfor index is later indexed
+    # with one then conservatively assume we can't fuse.
+    if index_positions is None:
+        index_positions = {}
+    if indexed_arrays is None:
+        indexed_arrays = set()
+    if non_indexed_arrays is None:
+        non_indexed_arrays = set()
+
+    def add_check_position(new_position,
+                           array_accessed,
+                           index_positions,
+                           indexed_arrays,
+                           non_indexed_arrays):
+        """Returns True if there is a reason to prevent fusion based
+           on the rules described above.
+           new_position will be a list or tuple of booleans that
+           says whether the index in that spot is a parfor index
+           or not.  array_accessed is the array on which the access
+           is occurring."""
+
+        # Convert list indices to tuple for generality.
+        if isinstance(new_position, list):
+            new_position = tuple(new_position)
+
+        # If none of the indices are based on a parfor index.
+        if True not in new_position:
+            # See if this array has been accessed before with a
+            # a parfor index and if so say that we can't fuse.
+            if array_accessed in indexed_arrays:
+                return True
+            else:
+                # Either array is already in non_indexed arrays or we
+                # will add it.  Either way, this index usage can fuse.
+                non_indexed_arrays.add(array_accessed)
+                return False
+
+        # Fallthrough for cases where one of the indices is a parfor index.
+        # If this array was previously accessed without a parfor index then
+        # conservatively say we can't fuse.
+        if array_accessed in non_indexed_arrays:
+            return True
+
+        npsize = len(new_position)
+        # See if we have not seen a npsize dimensioned array accessed before.
+        if npsize not in index_positions:
+            # If not then add current set of parfor/non-parfor indices and
+            # indicate it is safe as it is the first usage.
+            index_positions[npsize] = new_position
+            return False
+
+        # Here we have a subsequent access to a npsize-dimensioned array.
+        # Make sure we see the same combination of parfor/non-parfor index
+        # indices that we've seen before.  If not then return True saying
+        # that we can't fuse.
+        return index_positions[npsize] != new_position
+
+    def check_index(stmt_index,
+                    array_accessed,
+                    index_positions,
+                    indexed_arrays,
+                    non_indexed_arrays,
+                    derived_from_indices):
+        """Looks at the indices of a getitem or setitem to see if there
+           is a reason that they would prevent fusion.
+           Returns True if fusion should be prohibited, False otherwise.
+        """
+        if isinstance(stmt_index, ir.Var):
+            # If the array is 2+ dimensions then the index should be a tuple.
+            if isinstance(typemap[stmt_index.name], types.BaseTuple):
+                # Get how the index tuple is constructed.
+                fbs_res = guard(find_build_sequence, func_ir, stmt_index)
+                if fbs_res is not None:
+                    ind_seq, _ = fbs_res
+                    # If any indices are derived from an index is used then
+                    # return True to say we can't fuse.
+                    if (all([x.name in indices or
+                        x.name not in derived_from_indices for x in ind_seq])):
+                        # Get position in index tuple where parfor indices used.
+                        new_index_positions = [x.name in indices for x in ind_seq]
+                        # Make sure that we aren't accessing a given array with
+                        # different indices in a different order.
+                        return add_check_position(new_index_positions,
+                                                  array_accessed,
+                                                  index_positions,
+                                                  indexed_arrays,
+                                                  non_indexed_arrays)
+                    else:
+                        # index derived from a parfor index used so no fusion
+                        return True
+                else:
+                    # Don't know how the index tuple is built so
+                    # have to assume fusion can't happen.
+                    return True
+            else:
+                # Should be for 1D arrays.
+                if stmt_index.name in indices:
+                    # Array indexed by a parfor index variable.
+                    # Make sure this 1D access is consistent with prior ones.
+                    return add_check_position((True,),
+                                              array_accessed,
+                                              index_positions,
+                                              indexed_arrays,
+                                              non_indexed_arrays)
+                elif stmt_index.name in derived_from_indices:
+                    # If we ever index an array with something calculated
+                    # from an index then no fusion.
+                    return True
+                else:
+                    # Some kind of index that isn't a parfor index or
+                    # one derived from one, e.g., a constant.  Make sure
+                    # this is consistent with prior accessed of this array.
+                    return add_check_position((False,),
+                                              array_accessed,
+                                              index_positions,
+                                              indexed_arrays,
+                                              non_indexed_arrays)
+        else:
+            # We don't know how to handle non-Var indices so no fusion.
+            return True
+
+        # All branches above should cover all the cases and each should
+        # return so we should never get here.
+        raise errors.InternalError("Some code path in the parfor fusion "
+                                   "cross-iteration dependency checker "
+                                   "check_index didn't return a result.")
+
+    # Iterate through all the statements in the parfor.
     for b in parfor.loop_body.values():
         for stmt in b.body:
-            # GetItem/SetItem nodes are fine since can't have expression inside
-            # and only simple indices are possible
+            # Make sure SetItem accesses are fusion safe.
             if isinstance(stmt, (ir.SetItem, ir.StaticSetItem)):
+                if isinstance(typemap[stmt.target.name], types.npytypes.Array):
+                    # Check index safety with prior array accesses.
+                    if check_index(stmt.index,
+                                   stmt.target.name,
+                                   index_positions,
+                                   indexed_arrays,
+                                   non_indexed_arrays,
+                                   derived_from_indices):
+                        return True, index_positions, indexed_arrays, non_indexed_arrays
+                # Fusion safe so go to next statement.
                 continue
-            # tuples are immutable so no expression on parfor possible
-            if isinstance(stmt, ir.Assign) and isinstance(stmt.value, ir.Expr):
-                op = stmt.value.op
-                if op in ['build_tuple', 'getitem', 'static_getitem']:
-                    continue
-            # other statements can have potential violations
-            if not indices.isdisjoint(stmt.list_vars()):
-                dprint("has_cross_iter_dep found", indices, stmt)
-                return True
-    return False
+            elif isinstance(stmt, ir.Assign):
+                # If stmt of form a = parfor_index then add "a" to set of
+                # parfor indices.
+                if isinstance(stmt.value, ir.Var):
+                    if stmt.value.name in indices:
+                        indices.add(stmt.target.name)
+                        continue
+                elif isinstance(stmt.value, ir.Expr):
+                    op = stmt.value.op
+                    # Make sure getitem accesses are fusion safe.
+                    if op in ['getitem', 'static_getitem']:
+                        if isinstance(typemap[stmt.value.value.name], types.npytypes.Array):
+                            # Check index safety with prior array accesses.
+                            if check_index(stmt.value.index,
+                                           stmt.value.value.name,
+                                           index_positions,
+                                           indexed_arrays,
+                                           non_indexed_arrays,
+                                           derived_from_indices):
+                                return True, index_positions, indexed_arrays, non_indexed_arrays
+                        # Fusion safe so go to next statement.
+                        continue
+                    elif op == 'call':
+                        # If there is a call in the parfor body that takes some
+                        # array parameter then we have no way to analyze what
+                        # that call is doing so presume it is unsafe for fusion.
+                        if (any([isinstance(typemap[x.name], types.npytypes.Array)
+                            for x in stmt.value.list_vars()])):
+                            return True, index_positions, indexed_arrays, non_indexed_arrays
+
+                    # Get the vars used by this non-setitem/getitem statement.
+                    rhs_vars = [x.name for x in stmt.value.list_vars()]
+                    # If a parfor index is used as part of this statement or
+                    # something previous determined to be derived from a parfor
+                    # index then add the target variable to the set of
+                    # variables that are derived from parfors and so should
+                    # prevent fusion if used as an index.
+                    if (not indices.isdisjoint(rhs_vars) or
+                        not derived_from_indices.isdisjoint(rhs_vars)):
+                        derived_from_indices.add(stmt.target.name)
+
+    return False, index_positions, indexed_arrays, non_indexed_arrays
 
 
 def dprint(*s):


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
@@ -45,6 +45,7 @@
                       skip_parfors_unsupported, _32bit, needs_blas,
                       needs_lapack, disabled_test, skip_unless_scipy,
                       needs_subprocess)
+from numba.core.extending import register_jitable
 import cmath
 import unittest
 
@@ -1263,7 +1264,6 @@ def ddot(a, v):
                "hardware")
         self.assertIn(msg, str(raised.exception))
 
-
 @skip_parfors_unsupported
 class TestParfors(TestParforsBase):
     """ Tests cpython, reduction and various parfors features"""
@@ -1957,6 +1957,120 @@ def test_impl(x, sz):
             return x
         self.check(test_impl, np.zeros((10, 10)), 3)
 
+    def test_prange_unknown_call1(self):
+        @register_jitable
+        def issue7854_proc(u, i, even, size):
+            for j in range((even + i + 1) % 2 + 1, size - 1, 2):
+                u[i, j] = u[i + 1, j] + 1
+
+        # issue7854
+        # Forbid fusion in unanalyzable call inside prange.
+        def test_impl(u, size):
+            for i in numba.prange(1, size - 1):
+                issue7854_proc(u, i, 0, size)
+            for i in numba.prange(1, size - 1):
+                issue7854_proc(u, i, 1, size)
+            return u
+
+        size = 4
+        u = np.zeros((size, size))
+        cptypes = (numba.float64[:, ::1], types.int64)
+        self.assertEqual(countParfors(test_impl, cptypes), 2)
+        self.check(test_impl, u, size)
+
+    def test_prange_index_calc1(self):
+        # Should forbid fusion due to cross-iteration dependency as
+        # detected by loop index calcuation (i+1) as array index.
+        def test_impl(u, size):
+            for i in numba.prange(1, size - 1):
+                for j in range((i + 1) % 2 + 1, size - 1, 2):
+                    u[i, j] = u[i + 1, j] + 1
+            for i in numba.prange(1, size - 1):
+                for j in range(i % 2 + 1, size - 1, 2):
+                    u[i, j] = u[i + 1, j] + 1
+            return u
+
+        size = 4
+        u = np.zeros((size, size))
+        cptypes = (numba.float64[:, ::1], types.int64)
+        self.assertEqual(countParfors(test_impl, cptypes), 2)
+        self.check(test_impl, u, size)
+
+    def test_prange_reverse_order1(self):
+        # Testing if reversed loop index usage as array index
+        # prevents fusion.
+        def test_impl(a, b, size):
+            for i in numba.prange(size):
+                for j in range(size):
+                    a[i, j] = b[i, j] + 1
+            for i in numba.prange(size):
+                for j in range(size):
+                    b[j, i] = 3
+            return a[0, 0] + b[0, 0]
+
+        size = 10
+        a = np.zeros((size, size))
+        b = np.zeros((size, size))
+        cptypes = (numba.float64[:, ::1], numba.float64[:, ::1], types.int64)
+        self.assertEqual(countParfors(test_impl, cptypes), 2)
+        self.check(test_impl, a, b, size)
+
+    def test_prange_parfor_index_then_not(self):
+        # Testing if accessing an array first with a parfor index then
+        # without will prevent fusion.
+        def test_impl(a, size):
+            b = 0
+            for i in numba.prange(size):
+                a[i] = i
+            for i in numba.prange(size):
+                b += a[5]
+            return b
+
+        size = 10
+        a = np.zeros(size)
+        cptypes = (numba.float64[:], types.int64)
+        self.assertEqual(countParfors(test_impl, cptypes), 2)
+        self.check(test_impl, a, size)
+
+    def test_prange_parfor_index_const_tuple_fusion(self):
+        # Testing if accessing a tuple with prange index
+        # and later with a constant will not prevent fusion.
+        def test_impl(a, tup, size):
+            acc = 0
+            for i in numba.prange(size):
+                a[i] = i + tup[i]
+            for i in numba.prange(size):
+                acc += a[i] + tup[1]
+            return acc
+
+        size = 10
+        a = np.zeros(size)
+        b = tuple(a)
+        cptypes = (numba.float64[:],
+                   types.containers.UniTuple(types.float64, size),
+                   types.intp)
+        self.assertEqual(countParfors(test_impl, cptypes), 1)
+        self.check(test_impl, a, b, size)
+
+    def test_prange_non_parfor_index_then_opposite(self):
+        # Testing if accessing an array first without a parfor index then
+        # with will prevent fusion.
+        def test_impl(a, b, size):
+            for i in numba.prange(size):
+                b[i] = a[5]
+            for i in numba.prange(size):
+                a[i] = i
+            # Need this to stop previous prange from being optimized away.
+            b[0] += a[0]
+            return b
+
+        size = 10
+        a = np.zeros(size)
+        b = np.zeros(size)
+        cptypes = (numba.float64[:], numba.float64[:], types.int64)
+        self.assertEqual(countParfors(test_impl, cptypes), 2)
+        self.check(test_impl, a, b, size)
+
     def test_untraced_value_tuple(self):
         # This is a test for issue #6478.
         def test_impl():