Repository: numba/numba

Issue:
Segfault when using assert with parallel=True
The following code triggers a segfault. Tested on Ubuntu 22.04 with Numba 0.56.4 and Numba 0.57.0.

```python
from numba import njit
import numpy as np

@njit(parallel=True)
def foo():
    n = 1000

    for _ in range(n):
        for _ in range(n):
            values = np.zeros(3)

            np.max(values)

            assert n >= 0

foo()
```

- [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/main/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'.


Discussion:
This happens when `n` is too big. On my machine I can get `n` to be `323` and with `324` it segfaults.
ping @DrTodd13 
@99991 You can do the body of your nested for loops in another function and that should solve the problem.  I wonder if this MWR obscures the issue a bit because parallelizing a "max" over 3 elements isn't very useful and in fact will slow the program down.  Can the outer loops be collapsed into one loop with pndindex or one of the loops here made into a prange?  If so, then the "max" inside the loop nest will be lowered sequentially and the problem will also go away.  We think we know why this segfault is occurring and the fix is probably moderately difficult and we will be working on it but wanted to give some workarounds in the meantime.
> wanted to give some workarounds in the meantime.

Thank you!

>I wonder if this MWR obscures the issue a bit because parallelizing a "max" over 3 elements isn't very useful and in fact will slow the program down.

Definitely. But I doubt that anyone would read the 700 lines of filter kernels that this has been condensed from and removing the `prange` was required for the `M` in MWR.

> You can do the body of your nested for loops in another function and that should solve the problem.

Perhaps, but the original code would then need to pass about 16 parameters, which is a bit much I think.

> Can the outer loops be collapsed into one loop with pndindex or one of the loops here made into a prange? If so, then the "max" inside the loop nest will be lowered sequentially and the problem will also go away. 

The original code computes something for every pixel of an image without any pixel dependencies (embarrassingly parallel), so both of the loops could make use of `prange`. Hadn't heard of `pndindex` yet.

Anyway, when using `prange` for any loop and changing the assert to `i >= 0`, there are still segfaults:

```python
from numba import njit, prange
import numpy as np
import os

os.environ["NUMBA_PARALLEL_DIAGNOSTICS"] = "4"

@njit(parallel=True)
def foo():
    for y in prange(1000):
        for x in range(1000):
            values = np.zeros(3)

            i = np.argmax(values)

            assert i >= 0

foo()
```

I do not know how to read the diagnostics, but it seems that something is still happening with argmax.
@DrTodd13 Is it likely we can fix this by 0.58?
@99991 @stuartarchibald This doesn't appear to be documented anywhere but the parallelism code will only convert a prange to run in parallel if it is able to rewrite the loop into a parfor and it can only do that if there is only one exit from the loop.  Because you have an assert in there, there is more than 1 exit from the loop.  If you think about it, to some extent an assert here is ill-defined because normally you'd hit the first such extraneous control flow but in parallel you might hit a different assert or a different element would cause the assert.  I'm not sure this was planned but it is a bit fortunate because the parfor is turned into a gufunc and gufuncs swallow exceptions anyway.

So, if you take out the assert then the parfors in the loop body (which are the zeros and the argmax) are lowered sequentially and you don't get a segfault.  Suggestions Stuart?

Resolving patch:
diff --git a/numba/np/ufunc/parallel.py b/numba/np/ufunc/parallel.py
--- a/numba/np/ufunc/parallel.py
+++ b/numba/np/ufunc/parallel.py
@@ -515,6 +515,8 @@ def raise_with_hint(required):
             ll.add_symbol('numba_parallel_for', lib.parallel_for)
             ll.add_symbol('do_scheduling_signed', lib.do_scheduling_signed)
             ll.add_symbol('do_scheduling_unsigned', lib.do_scheduling_unsigned)
+            ll.add_symbol('allocate_sched', lib.allocate_sched)
+            ll.add_symbol('deallocate_sched', lib.deallocate_sched)
 
             launch_threads = CFUNCTYPE(None, c_int)(lib.launch_threads)
             launch_threads(NUM_THREADS)
diff --git a/numba/parfors/parfor.py b/numba/parfors/parfor.py
--- a/numba/parfors/parfor.py
+++ b/numba/parfors/parfor.py
@@ -24,6 +24,7 @@
 from contextlib import contextmanager
 import operator
 from dataclasses import make_dataclass
+import warnings
 
 from llvmlite import ir as lir
 from numba.core.imputils import impl_ret_untracked
@@ -2359,7 +2360,27 @@ def run(self, blocks):
         # We go over all loops, smaller loops first (inner first)
         for loop, s in sorted(sized_loops, key=lambda tup: tup[1]):
             if len(loop.entries) != 1 or len(loop.exits) != 1:
+                if not config.DISABLE_PERFORMANCE_WARNINGS:
+                    for entry in loop.entries:
+                        for inst in blocks[entry].body:
+                            # if prange or pndindex call
+                            if (
+                                isinstance(inst, ir.Assign)
+                                and isinstance(inst.value, ir.Expr)
+                                and inst.value.op == "call"
+                                and self._is_parallel_loop(
+                                    inst.value.func.name, call_table)
+                            ):
+                                msg = "\nprange or pndindex loop " \
+                                      "will not be executed in " \
+                                      "parallel due to there being more than one " \
+                                      "entry to or exit from the loop (e.g., an " \
+                                      "assertion)."
+                                warnings.warn(
+                                    errors.NumbaPerformanceWarning(
+                                        msg, inst.loc))
                 continue
+
             entry = list(loop.entries)[0]
             for inst in blocks[entry].body:
                 # if prange or pndindex call
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
@@ -1668,7 +1668,9 @@ def load_range(v):
     intp_t = context.get_value_type(types.intp)
     uintp_t = context.get_value_type(types.uintp)
     intp_ptr_t = llvmlite.ir.PointerType(intp_t)
+    intp_ptr_ptr_t = llvmlite.ir.PointerType(intp_ptr_t)
     uintp_ptr_t = llvmlite.ir.PointerType(uintp_t)
+    uintp_ptr_ptr_t = llvmlite.ir.PointerType(uintp_ptr_t)
     zero = context.get_constant(types.uintp, 0)
     one = context.get_constant(types.uintp, 1)
     one_type = one.type
@@ -1684,9 +1686,11 @@ def load_range(v):
     if index_var_typ.signed:
         sched_type = intp_t
         sched_ptr_type = intp_ptr_t
+        sched_ptr_ptr_type = intp_ptr_ptr_t
     else:
         sched_type = uintp_t
         sched_ptr_type = uintp_ptr_t
+        sched_ptr_ptr_type = uintp_ptr_ptr_t
 
     # Call do_scheduling with appropriate arguments
     dim_starts = cgutils.alloca_once(
@@ -1713,6 +1717,7 @@ def load_range(v):
         builder.store(stop, builder.gep(dim_stops,
                                         [context.get_constant(types.uintp, i)]))
 
+    # Prepare to call get/set parallel_chunksize and get the number of threads.
     get_chunksize = cgutils.get_or_insert_function(
         builder.module,
         llvmlite.ir.FunctionType(uintp_t, []),
@@ -1728,7 +1733,9 @@ def load_range(v):
         llvmlite.ir.FunctionType(llvmlite.ir.IntType(types.intp.bitwidth), []),
         "get_num_threads")
 
+    # Get the current number of threads.
     num_threads = builder.call(get_num_threads, [])
+    # Get the current chunksize so we can use it and restore the value later.
     current_chunksize = builder.call(get_chunksize, [])
 
     with cgutils.if_unlikely(builder, builder.icmp_signed('<=', num_threads,
@@ -1738,6 +1745,9 @@ def load_range(v):
                                                   ("Invalid number of threads. "
                                                    "This likely indicates a bug in Numba.",))
 
+    # Call get_sched_size from gufunc_scheduler.cpp that incorporates the size of the work,
+    # the number of threads and the selected chunk size.  This will tell us how many entries
+    # in the schedule we will need.
     get_sched_size_fnty = llvmlite.ir.FunctionType(uintp_t, [uintp_t, uintp_t, intp_ptr_t, intp_ptr_t])
     get_sched_size = cgutils.get_or_insert_function(
         builder.module,
@@ -1747,11 +1757,27 @@ def load_range(v):
                                                   context.get_constant(types.uintp, num_dim),
                                                   dim_starts,
                                                   dim_stops])
+    # Set the chunksize to zero so that any nested calls get the default chunk size behavior.
     builder.call(set_chunksize, [zero])
 
+    # Each entry in the schedule is 2 times the number of dimensions long.
     multiplier = context.get_constant(types.uintp, num_dim * 2)
+    # Compute the total number of entries in the schedule.
     sched_size = builder.mul(num_divisions, multiplier)
-    sched = builder.alloca(sched_type, size=sched_size, name="sched")
+
+    # Prepare to dynamically allocate memory to hold the schedule.
+    alloc_sched_fnty = llvmlite.ir.FunctionType(sched_ptr_type, [uintp_t])
+    alloc_sched_func = cgutils.get_or_insert_function(
+        builder.module,
+        alloc_sched_fnty,
+        name="allocate_sched")
+    # Call gufunc_scheduler.cpp to allocate the schedule.
+    # This may or may not do pooling.
+    alloc_space = builder.call(alloc_sched_func, [sched_size])
+    # Allocate a slot in the entry block to store the schedule pointer.
+    sched = cgutils.alloca_once(builder, sched_ptr_type)
+    # Store the schedule pointer into that slot.
+    builder.store(alloc_space, sched)
 
     debug_flag = 1 if config.DEBUG_ARRAY_OPT else 0
     scheduling_fnty = llvmlite.ir.FunctionType(
@@ -1765,11 +1791,12 @@ def load_range(v):
                                                        scheduling_fnty,
                                                        name="do_scheduling_unsigned")
 
+    # Call the scheduling routine that decides how to break up the work.
     builder.call(
         do_scheduling, [
             context.get_constant(
                 types.uintp, num_dim), dim_starts, dim_stops, num_divisions,
-            sched, context.get_constant(
+            builder.load(sched), context.get_constant(
                     types.intp, debug_flag)])
 
     # Get the LLVM vars for the Numba IR reduction array vars.
@@ -1810,7 +1837,7 @@ def load_potential_tuple_var(x):
         name="pargs")
     array_strides = []
     # sched goes first
-    builder.store(builder.bitcast(sched, byte_ptr_t), args)
+    builder.store(builder.bitcast(builder.load(sched), byte_ptr_t), args)
     array_strides.append(context.get_constant(types.intp, sizeof_intp))
     rv_to_arg_dict = {}
     # followed by other arguments
@@ -1964,6 +1991,14 @@ def load_potential_tuple_var(x):
 
     builder.call(set_chunksize, [current_chunksize])
 
+    # Deallocate the schedule's memory.
+    dealloc_sched_fnty = llvmlite.ir.FunctionType(llvmlite.ir.VoidType(), [sched_ptr_type])
+    dealloc_sched_func = cgutils.get_or_insert_function(
+        builder.module,
+        dealloc_sched_fnty,
+        name="deallocate_sched")
+    builder.call(dealloc_sched_func, [builder.load(sched)])
+
     for k, v in rv_to_arg_dict.items():
         arg, rv_arg = v
         only_elem_ptr = builder.gep(rv_arg, [context.get_constant(types.intp, 0)])


Test patch:
diff --git a/numba/tests/test_debug.py b/numba/tests/test_debug.py
--- a/numba/tests/test_debug.py
+++ b/numba/tests/test_debug.py
@@ -46,6 +46,13 @@ def supported_parfor(n):
         a[i] = a[i] + np.sin(i)
     return a
 
+def unsupported_prange(n):
+    a = np.ones(n)
+    for i in prange(n):
+        a[i] = a[i] + np.sin(i)
+        assert i + 13 < 100000
+    return a
+
 force_parallel_flags = Flags()
 force_parallel_flags.auto_parallel = ParallelOptions(True)
 force_parallel_flags.nrt = True
@@ -241,6 +248,17 @@ def check_parfors_warning(self, warn_list):
                 break
         self.assertTrue(warning_found, "Warning message should be found.")
 
+    def check_parfors_unsupported_prange_warning(self, warn_list):
+        msg = ("prange or pndindex loop will not be executed in parallel "
+               "due to there being more than one entry to or exit from the "
+               "loop (e.g., an assertion).")
+        warning_found = False
+        for w in warn_list:
+            if msg in str(w.message):
+                warning_found = True
+                break
+        self.assertTrue(warning_found, "Warning message should be found.")
+
     @needs_blas
     @skip_parfors_unsupported
     def test_warns(self):
@@ -255,6 +273,18 @@ def test_warns(self):
                                     flags=force_parallel_flags)
         self.check_parfors_warning(w)
 
+    @needs_blas
+    @skip_parfors_unsupported
+    def test_unsupported_prange_warns(self):
+        """
+        Test that prange with multiple exits issues a warning
+        """
+        with warnings.catch_warnings(record=True) as w:
+            warnings.simplefilter("always", NumbaPerformanceWarning)
+            cres = compile_isolated(unsupported_prange, (types.int64,),
+                                    flags=force_parallel_flags)
+        self.check_parfors_unsupported_prange_warning(w)
+
     @skip_parfors_unsupported
     def test_array_debug_opt_stats(self):
         """
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
@@ -2285,6 +2285,33 @@ def test_impl(n):
         self.check(test_impl, 15)
         self.assertEqual(countParfors(test_impl, (types.int64, )), 2)
 
+    def test_issue9029(self):
+        # issue9029: too many parfors executed in one function
+        # overflowed the stack.
+        def test_impl(i1, i2):
+            N = 30
+            S = 3
+            a = np.empty((N,N))
+            # The stack should overflow if there are 30*30*2 (# of parfors)
+            # iterations.
+            for y in range(N):
+                for x in range(N):
+                    values = np.ones(S)
+                    v = values[0]
+
+                    p2 = np.empty(S)
+                    for i in prange(i1, i2):
+                        p2[i] = 1
+                    j = p2[0]
+
+                    a[y,x] = v + j
+            return a
+
+        # We pass in 0 and 3 so that the function can't analyze the loop
+        # bounds on the prange to generate a signed loop whereas the
+        # np.ones will be an unsigned loop.
+        self.check(test_impl, 0, 3)
+
     def test_fusion_no_side_effects(self):
         def test_impl(a, b):
             X = np.ones(100)