Repository: numba/numba

Issue:
Truediv Parfor fails when input array and return value has control flow
<!--

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/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'.

<!--

Please include details of the bug here, including, if applicable, what you
expected to happen!

-->

Here is a dummy reproducer to show the issue.

```
import numpy as np
import numba

@numba.njit(parallel=True)
def f(arr, flag1, flag2):  # pragma: no cover
    inv = np.arange(arr.size)
    if flag1:
        return inv.astype(np.float64)
    if flag2:
        ret = inv[inv]
    else:
        ret = inv[inv - 1]
    return ret / arr.size

arr = np.arange(100)
print(f(arr, True, False))
```

This outputs:

```python3
func_fnty: Function(<function runtime_broadcast_assert_shapes at 0x1154f9700>)
typs: [array(int64, 1d, C)]
dims: [1] 1
fargs: [Literal[int](1), array(int64, 1d, C)]
Traceback (most recent call last):
  File "/Users/nicholasriasanovsky/Documents/bodo/Bodo/test_numba.py", line 16, in <module>
    print(f(arr, True, False))
  File "/Users/nicholasriasanovsky/Documents/bodo/numba/numba/core/dispatcher.py", line 468, in _compile_for_args
    error_rewrite(e, 'typing')
  File "/Users/nicholasriasanovsky/Documents/bodo/numba/numba/core/dispatcher.py", line 407, in error_rewrite
    raise e
  File "/Users/nicholasriasanovsky/Documents/bodo/numba/numba/core/dispatcher.py", line 420, in _compile_for_args
    return_val = self.compile(tuple(argtypes))
  File "/Users/nicholasriasanovsky/Documents/bodo/numba/numba/core/dispatcher.py", line 965, in compile
    cres = self._compiler.compile(args, return_type)
  File "/Users/nicholasriasanovsky/Documents/bodo/numba/numba/core/dispatcher.py", line 129, in compile
    raise retval
  File "/Users/nicholasriasanovsky/Documents/bodo/numba/numba/core/dispatcher.py", line 139, in _compile_cached
    retval = self._compile_core(args, return_type)
  File "/Users/nicholasriasanovsky/Documents/bodo/numba/numba/core/dispatcher.py", line 152, in _compile_core
    cres = compiler.compile_extra(self.targetdescr.typing_context,
  File "/Users/nicholasriasanovsky/Documents/bodo/numba/numba/core/compiler.py", line 716, in compile_extra
    return pipeline.compile_extra(func)
  File "/Users/nicholasriasanovsky/Documents/bodo/numba/numba/core/compiler.py", line 452, in compile_extra
    return self._compile_bytecode()
  File "/Users/nicholasriasanovsky/Documents/bodo/numba/numba/core/compiler.py", line 520, in _compile_bytecode
    return self._compile_core()
  File "/Users/nicholasriasanovsky/Documents/bodo/numba/numba/core/compiler.py", line 499, in _compile_core
    raise e
  File "/Users/nicholasriasanovsky/Documents/bodo/numba/numba/core/compiler.py", line 486, in _compile_core
    pm.run(self.state)
  File "/Users/nicholasriasanovsky/Documents/bodo/numba/numba/core/compiler_machinery.py", line 368, in run
    raise patched_exception
  File "/Users/nicholasriasanovsky/Documents/bodo/numba/numba/core/compiler_machinery.py", line 356, in run
    self._runPass(idx, pass_inst, state)
  File "/Users/nicholasriasanovsky/Documents/bodo/numba/numba/core/compiler_lock.py", line 35, in _acquire_compile_lock
    return func(*args, **kwargs)
  File "/Users/nicholasriasanovsky/Documents/bodo/numba/numba/core/compiler_machinery.py", line 311, in _runPass
    mutated |= check(pss.run_pass, internal_state)
  File "/Users/nicholasriasanovsky/Documents/bodo/numba/numba/core/compiler_machinery.py", line 273, in check
    mangled = func(compiler_state)
  File "/Users/nicholasriasanovsky/Documents/bodo/numba/numba/core/typed_passes.py", line 306, in run_pass
    parfor_pass.run()
  File "/Users/nicholasriasanovsky/Documents/bodo/numba/numba/parfors/parfor.py", line 2864, in run
    self._pre_run()
  File "/Users/nicholasriasanovsky/Documents/bodo/numba/numba/parfors/parfor.py", line 2856, in _pre_run
    self.array_analysis.run(self.func_ir.blocks)
  File "/Users/nicholasriasanovsky/Documents/bodo/numba/numba/parfors/array_analysis.py", line 1178, in run
    self._run_on_blocks(topo_order, blocks, cfg, init_equiv_set)
  File "/Users/nicholasriasanovsky/Documents/bodo/numba/numba/parfors/array_analysis.py", line 1198, in _run_on_blocks
    pending_transforms = self._determine_transform(
  File "/Users/nicholasriasanovsky/Documents/bodo/numba/numba/parfors/array_analysis.py", line 1273, in _determine_transform
    pre, post = self._analyze_inst(
  File "/Users/nicholasriasanovsky/Documents/bodo/numba/numba/parfors/array_analysis.py", line 1311, in _analyze_inst
    result = self._analyze_expr(scope, equiv_set, inst.value, lhs)
  File "/Users/nicholasriasanovsky/Documents/bodo/numba/numba/parfors/array_analysis.py", line 1571, in _analyze_expr
    return guard(fn, scope, equiv_set, expr, lhs)
  File "/Users/nicholasriasanovsky/Documents/bodo/numba/numba/core/ir_utils.py", line 1527, in guard
    return func(*args, **kwargs)
  File "/Users/nicholasriasanovsky/Documents/bodo/numba/numba/parfors/array_analysis.py", line 2183, in _analyze_op_arrayexpr
    return self._analyze_broadcast(
  File "/Users/nicholasriasanovsky/Documents/bodo/numba/numba/parfors/array_analysis.py", line 3056, in _analyze_broadcast
    return self._insert_runtime_broadcast_call(
  File "/Users/nicholasriasanovsky/Documents/bodo/numba/numba/parfors/array_analysis.py", line 2982, in _insert_runtime_broadcast_call
    func_sig = self.context.resolve_function_type(func_fnty, fargs, {})
  File "/Users/nicholasriasanovsky/Documents/bodo/numba/numba/core/typing/context.py", line 213, in resolve_function_type
    raise last_exception
  File "/Users/nicholasriasanovsky/Documents/bodo/numba/numba/core/typing/context.py", line 196, in resolve_function_type
    res = self._resolve_user_function_type(func, args, kws)
  File "/Users/nicholasriasanovsky/Documents/bodo/numba/numba/core/typing/context.py", line 248, in _resolve_user_function_type
    return func.get_call_type(self, args, kws)
  File "/Users/nicholasriasanovsky/Documents/bodo/numba/numba/core/types/functions.py", line 331, in get_call_type
    failures.raise_error()
  File "/Users/nicholasriasanovsky/Documents/bodo/numba/numba/core/types/functions.py", line 227, in raise_error
    raise errors.TypingError(self.format())
numba.core.errors.TypingError: Failed in nopython mode pipeline (step: convert to parfors)
No implementation of function Function(<function runtime_broadcast_assert_shapes at 0x1154f9700>) found for signature:
 
 >>> runtime_broadcast_assert_shapes(Literal[int](1), array(int64, 1d, C))
 
There are 2 candidate implementations:
 - Of which 2 did not match due to:
 Overload in function 'register_jitable.<locals>.wrap.<locals>.ov_wrap': File: numba/core/extending.py: Line 159.
   With argument(s): '(int64, array(int64, 1d, C))':
  Rejected as the implementation raised a specific error:
    TypeError: missing a required argument: 'arg1'
      Traceback (most recent call last):
        File "/Users/nicholasriasanovsky/Documents/bodo/numba/numba/core/types/functions.py", line 308, in get_call_type
          sig = temp.apply(nolitargs, nolitkws)
        File "/Users/nicholasriasanovsky/Documents/bodo/numba/numba/core/typing/templates.py", line 351, in apply
          sig = generic(args, kws)
        File "/Users/nicholasriasanovsky/Documents/bodo/numba/numba/core/typing/templates.py", line 614, in generic
          disp, new_args = self._get_impl(args, kws)
        File "/Users/nicholasriasanovsky/Documents/bodo/numba/numba/core/typing/templates.py", line 713, in _get_impl
          impl, args = self._build_impl(cache_key, args, kws)
        File "/Users/nicholasriasanovsky/Documents/bodo/numba/numba/core/typing/templates.py", line 824, in _build_impl
          disp_type.get_call_type(self.context, args, kws)
        File "/Users/nicholasriasanovsky/Documents/bodo/numba/numba/core/types/functions.py", line 541, in get_call_type
          self.dispatcher.get_call_template(args, kws)
        File "/Users/nicholasriasanovsky/Documents/bodo/numba/numba/core/dispatcher.py", line 359, in get_call_template
          pysig, args = self._compiler.fold_argument_types(args, kws)
        File "/Users/nicholasriasanovsky/Documents/bodo/numba/numba/core/dispatcher.py", line 118, in fold_argument_types
          args = fold_arguments(self.pysig, args, kws,
        File "/Users/nicholasriasanovsky/Documents/bodo/numba/numba/core/typing/templates.py", line 221, in fold_arguments
          ba = pysig.bind(*bind_args, **bind_kws)
        File "/Users/nicholasriasanovsky/mambaforge/envs/DEV/lib/python3.9/inspect.py", line 3050, in bind
          return self._bind(args, kwargs)
        File "/Users/nicholasriasanovsky/mambaforge/envs/DEV/lib/python3.9/inspect.py", line 2965, in _bind
          raise TypeError(msg) from None
      TypeError: missing a required argument: 'arg1'
  raised from /Users/nicholasriasanovsky/mambaforge/envs/DEV/lib/python3.9/inspect.py:2965
```

which occurs because `runtime_broadcast_assert_shapes` tries to run with 1 array. Skipping the `_insert_runtime_broadcast_call` by returning None produces. 

```python3
numba.core.errors.LoweringError: Failed in nopython mode pipeline (step: native lowering)
'$62binary_true_divide.3_size0.11'

File "test_numba.py", line 13:
def f(arr, flag1, flag2):  # pragma: no cover
    <source elided>
        ret = inv[inv - 1]
    return ret / arr.size
```
because the parfor uses `$62binary_true_divide.3_size0.11` as an index but it is not defined in the IR.



Discussion:

Resolving patch:
diff --git a/numba/np/unsafe/ndarray.py b/numba/np/unsafe/ndarray.py
--- a/numba/np/unsafe/ndarray.py
+++ b/numba/np/unsafe/ndarray.py
@@ -77,4 +77,3 @@ def impl(array, length, empty_tuple):
         return res
 
     return sig, codegen
-
diff --git a/numba/parfors/array_analysis.py b/numba/parfors/array_analysis.py
--- a/numba/parfors/array_analysis.py
+++ b/numba/parfors/array_analysis.py
@@ -26,9 +26,8 @@
 from numba.core.analysis import compute_cfg_from_blocks
 from numba.core.typing import npydecl, signature
 import copy
-from numba.core.extending import intrinsic, register_jitable
+from numba.core.extending import intrinsic
 import llvmlite
-from numba.np.unsafe.ndarray import to_fixed_tuple
 
 UNKNOWN_CLASS = -1
 CONST_CLASS = 0
@@ -1074,27 +1073,6 @@ def __init__(self, slice_size, dim_size):
         self.dim_size = dim_size
 
 
-@register_jitable
-def runtime_broadcast_assert_shapes(max_dim, arg0, arg1):
-    new_shape = []
-    shapes = []
-    shapes.append(list(arg0.shape))
-    shapes.append(list(arg1.shape))
-    for i in range(max_dim):
-        sizes = []
-        for shape in shapes:
-            if i < len(shape):
-                size = shape[len(shape) - 1 - i]
-                if size != 1:
-                    sizes.append(size)  # non-1 size to front
-        if len(sizes) == 0:
-            sizes.append(1)
-        new_shape.append(sizes[0])
-
-    rev = new_shape[::-1]
-    return to_fixed_tuple(rev, max_dim)
-
-
 class ArrayAnalysis(object):
     aa_count = 0
 
@@ -1379,7 +1357,7 @@ def _analyze_inst(self, label, scope, equiv_set, inst, redefined):
                     shape is not None
                     and isinstance(shape, ir.Var)
                     and isinstance(
-                        self.typemap[shape.name], types.containers.Tuple
+                        self.typemap[shape.name], types.containers.BaseTuple
                     )
                 ):
                     pass
@@ -2950,58 +2928,6 @@ def _analyze_op_call_numpy_linalg_inv(
         require(len(args) >= 1)
         return ArrayAnalysis.AnalyzeResult(shape=equiv_set._get_shape(args[0]))
 
-    def _insert_runtime_broadcast_call(self, scope, loc, arrs, max_dim):
-        pre = []
-
-        typs = [self.typemap[x.name] for x in arrs]
-        dims = [self.typemap[x.name].ndim for x in arrs]
-        max_dim = max(dims)
-
-        runtime_broadcast_shape = ir.Var(
-            scope, mk_unique_var("runtime_broadcast_shape"), loc
-        )
-        runtime_broadcast_type = types.containers.List(types.intp)
-        self.typemap[runtime_broadcast_shape.name] = runtime_broadcast_type
-
-        func_var = ir.Var(scope, mk_unique_var("runtime_broadcast_call"), loc)
-        func_def = ir.Global(
-            "runtime_broadcast_assert_shapes",
-            runtime_broadcast_assert_shapes,
-            loc=loc
-        )
-        pre.append(
-            ir.Assign(value=func_def, target=func_var, loc=loc)
-        )
-        func_fnty = get_global_func_typ(runtime_broadcast_assert_shapes)
-        print("func_fnty:", func_fnty)
-        self.typemap[func_var.name] = func_fnty
-        fargs = [types.literal(max_dim)] + typs
-        print("typs:", typs)
-        print("dims:", dims, max_dim)
-        print("fargs:", fargs)
-        func_sig = self.context.resolve_function_type(func_fnty, fargs, {})
-        print("func_sig:", func_sig)
-
-        max_dim_var = ir.Var(scope, mk_unique_var("max_dim"), loc)
-        max_dim_val = ir.Const(max_dim, loc)
-        max_dim_typ = types.IntegerLiteral(max_dim)
-        self.typemap[max_dim_var.name] = max_dim_typ
-        pre.append(
-            ir.Assign(value=max_dim_val, target=max_dim_var, loc=loc)
-        )
-
-        func_call = ir.Expr.call(func_var, [max_dim_var] + arrs, {}, loc)
-
-        self.calltypes[func_call] = func_sig
-        pre.append(
-            ir.Assign(value=func_call, target=runtime_broadcast_shape, loc=loc)
-        )
-
-        return ArrayAnalysis.AnalyzeResult(
-            shape=runtime_broadcast_shape,
-            pre=pre
-        )
-
     def _analyze_broadcast(self, scope, equiv_set, loc, args, fn):
         """Infer shape equivalence of arguments based on Numpy broadcast rules
         and return shape of output
@@ -3048,14 +2974,35 @@ def _analyze_broadcast(self, scope, equiv_set, loc, args, fn):
                 shape=arrs[0],
                 pre=self._call_assert_equiv(scope, loc, equiv_set, arrs)
             )
-        if None not in shapes:
-            return self._broadcast_assert_shapes(
-                scope, equiv_set, loc, shapes, names
-            )
-        else:
-            return self._insert_runtime_broadcast_call(
-                scope, loc, arrs, max_dim
-            )
+        pre = []
+        if None in shapes:
+            # There is at least 1 shape that we don't know,
+            # so we need to generate that shape now.
+            new_shapes = []
+            for i, s in enumerate(shapes):
+                if s is None:
+                    var = arrs[i]
+                    typ = self.typemap[var.name]
+                    shape = self._gen_shape_call(
+                        equiv_set, var, typ.ndim, None, pre
+                    )
+                    new_shapes.append(shape)
+                else:
+                    new_shapes.append(s)
+            shapes = new_shapes
+
+        result = self._broadcast_assert_shapes(
+            scope, equiv_set, loc, shapes, names
+        )
+        if pre:
+            # If we had to generate a shape we have to insert
+            # that code before the broadcast assertion.
+            if 'pre' in result.kwargs:
+                prev_pre = result.kwargs['pre']
+            else:
+                prev_pre = []
+            result.kwargs['pre'] = pre + prev_pre
+        return result
 
     def _broadcast_assert_shapes(self, scope, equiv_set, loc, shapes, names):
         """Produce assert_equiv for sizes in each dimension, taking into


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
@@ -2059,6 +2059,55 @@ def test_impl(x):
 
         self.check(test_impl, np.zeros((3, 1)))
 
+    def test_1array_control_flow(self):
+        # issue8146
+        def test_impl(arr, flag1, flag2):
+            inv = np.arange(arr.size)
+            if flag1:
+                return inv.astype(np.float64)
+            if flag2:
+                ret = inv[inv]
+            else:
+                ret = inv[inv - 1]
+            return ret / arr.size
+
+        arr = np.arange(100)
+        self.check(test_impl, arr, True, False)
+        self.check(test_impl, arr, True, True)
+        self.check(test_impl, arr, False, False)
+
+    def test_2array_1_control_flow(self):
+        # issue8146
+        def test_impl(arr, l, flag):
+            inv1 = np.arange(arr.size)
+            inv2 = np.arange(l, arr.size + l)
+            if flag:
+                ret = inv1[inv1]
+            else:
+                ret = inv1[inv1 - 1]
+            return ret / inv2
+
+        arr = np.arange(100)
+        self.check(test_impl, arr, 10, True)
+        self.check(test_impl, arr, 10, False)
+
+    def test_2array_2_control_flow(self):
+        # issue8146
+        def test_impl(arr, l, flag):
+            inv1 = np.arange(arr.size)
+            inv2 = np.arange(l, arr.size + l)
+            if flag:
+                ret1 = inv1[inv1]
+                ret2 = inv2[inv1]
+            else:
+                ret1 = inv1[inv1 - 1]
+                ret2 = inv2[inv1 - 1]
+            return ret1 / ret2
+
+        arr = np.arange(100)
+        self.check(test_impl, arr, 10, True)
+        self.check(test_impl, arr, 10, False)
+
 @skip_parfors_unsupported
 class TestParforsLeaks(MemoryLeakMixin, TestParforsBase):
     def check(self, pyfunc, *args, **kwargs):