Repository: numba/numba

Issue:
Operations over a custom object aren't fused
- [ ] 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).

```python
import numpy
import numba
from numba import cgutils, types
from numba.extending import (box, intrinsic, make_attribute_wrapper, models,
                             overload_method, overload, register_model)
from numba.targets.boxing import box_array
from numba.typing import signature


class Dummy:
    def __init__(self, data):
        self.data = data

    def floordiv(self, other):
        return data // other.data

    def pow(self, other):
        return data ** other.data


class DummyType(types.Type):
    def __init__(self, data):
        self.data = data
        super().__init__(name=f'Dummy({self.data})')


@register_model(DummyType)
class DummyTypeModel(models.StructModel):
    def __init__(self, dmm, fe_type):
        members = [
            ('data', fe_type.data),
        ]
        super().__init__(dmm, fe_type, members)


make_attribute_wrapper(DummyType, 'data', '_data')


@intrinsic
def init_dummy(typingctx, data):
    def codegen(context, builder, sig, args):
        typ = sig.return_type
        data, = args
        dummy = cgutils.create_struct_proxy(typ)(context, builder)
        dummy.data = data

        if context.enable_nrt:
            context.nrt.incref(builder, sig.args[0], data)

        return dummy._getvalue()

    ret_typ = DummyType(data)
    sig = signature(ret_typ, data)

    return sig, codegen


@overload(Dummy, inline='always')
def dummy_overload(data):
    def ctor(data):
        return init_dummy(data)

    return ctor


@box(DummyType)
def box_dummy(typ, val, c):
    dummy = cgutils.create_struct_proxy(typ)(c.context, c.builder, value=val)
    data_obj = box_array(typ.data, dummy.data, c)
    class_obj = c.pyapi.unserialize(c.pyapi.serialize_object(Dummy))
    res = c.pyapi.call_function_objargs(class_obj, (data_obj,))
    c.pyapi.decref(data_obj)
    c.pyapi.decref(class_obj)

    return res


@overload_method(DummyType, 'floordiv', inline='always')
def floordiv_overload(self, other):
    def impl(self, other):
        return self._data // other._data

    return impl


@overload_method(DummyType, 'pow', inline='always')
def pow_overload(self, other):
    def impl(self, other):
        return self._data ** other._data

    return impl


@numba.njit(parallel=True)
def test_impl1(data, other_data):
    dummy = Dummy(data)
    other = Dummy(other_data)

    # After fusion, function test_impl1 has 2 parallel for-loop(s) #{0, 1}.
    res1 = dummy.floordiv(other)
    res2 = dummy.pow(other)

    return res1, res2


@numba.njit(parallel=True)
def test_impl2(data, other_data):
    # Parallel for-loop #3 is fused into for-loop #2.
    # After fusion, function test_impl2 has 1 parallel for-loop(s) #{2}.
    res1 = data // other_data
    res2 = data ** other_data

    return res1, res2


data = numpy.arange(10)
other_data = numpy.array([1, 2] * 5)
# not fused
res1, res2 = test_impl1(data, other_data)
# fused
res3, res4 = test_impl2(data, other_data)
```

Operations are not fused in `test_impl1`, but fused in `test_impl2`. In the first case the operations on arrays are performed inside of the custom object `Dummy`. Am I able to fuse such operations somehow?

Numba version: ef119bc



Discussion:
@densmirn thanks for submitting this to the Numba issue tracker. Would it be possible to add some instructions for  showing if operations are fused or are not fused? This would enable me to attempt to reproduce.
@esc look at the dispatcher `.parallel_diagnostics` method.
@stuartarchibald thanks, here is what I see:

```
In [2]: test_impl1.parallel_diagnostics()

================================================================================
 Parallel Accelerator Optimizing:  Function test_impl1,
/Users/vhaenel/git/numba/foo031.py (94)
================================================================================


Parallel loop listing for  Function test_impl1, /Users/vhaenel/git/numba/foo031.py (94)
-------------------------------------------------------------------------------|loop #ID
@numba.njit(parallel=True)                                                     |
def test_impl1(data, other_data):                                              |
    dummy = Dummy(data)                                                        |
    other = Dummy(other_data)                                                  |
                                                                               |
    # After fusion, function test_impl1 has 2 parallel for-loop(s) #{0, 1}.    |
    res1 = dummy.floordiv(other)                                               |
    res2 = dummy.pow(other)                                                    |
                                                                               |
    return res1, res2                                                          |
------------------------------ After Optimisation ------------------------------
Parallel structure is already optimal.
--------------------------------------------------------------------------------
--------------------------------------------------------------------------------


In [3]: test_impl2.parallel_diagnostics()

================================================================================
 Parallel Accelerator Optimizing:  Function test_impl2,
/Users/vhaenel/git/numba/foo031.py (106)
================================================================================


Parallel loop listing for  Function test_impl2, /Users/vhaenel/git/numba/foo031.py (106)
----------------------------------------------------------------------------|loop #ID
@numba.njit(parallel=True)                                                  |
def test_impl2(data, other_data):                                           |
    # Parallel for-loop #3 is fused into for-loop #2.                       |
    # After fusion, function test_impl2 has 1 parallel for-loop(s) #{2}.    |
    res1 = data // other_data-----------------------------------------------| #2
    res2 = data ** other_data-----------------------------------------------| #3
                                                                            |
    return res1, res2                                                       |
------------------------------ After Optimisation ------------------------------
Parallel region 0:
+--2 (parallel, fused with loop(s): 3)



Parallel region 0 (loop #2) had 1 loop(s) fused.
--------------------------------------------------------------------------------
--------------------------------------------------------------------------------
```
Perhaps turn the level up a bit http://numba.pydata.org/numba-doc/dev/user/parallel.html#diagnostics
Here is what I see with an increased level:

```
In [2]: test_impl1.parallel_diagnostics(level=4)

================================================================================
 Parallel Accelerator Optimizing:  Function test_impl1,
/Users/vhaenel/git/numba/foo031.py (94)
================================================================================


Parallel loop listing for  Function test_impl1, /Users/vhaenel/git/numba/foo031.py (94)
-------------------------------------------------------------------------------|loop #ID
@numba.njit(parallel=True)                                                     |
def test_impl1(data, other_data):                                              |
    dummy = Dummy(data)                                                        |
    other = Dummy(other_data)                                                  |
                                                                               |
    # After fusion, function test_impl1 has 2 parallel for-loop(s) #{0, 1}.    |
    res1 = dummy.floordiv(other)                                               |
    res2 = dummy.pow(other)                                                    |
                                                                               |
    return res1, res2                                                          |
--------------------------------- Fusing loops ---------------------------------
Attempting fusion of parallel loops (combines loops with similar properties)...
  Trying to fuse loops #0 and #1:
    - fusion failed: loop dimension mismatched in axis 0. slice(0,
$4load_attr.1.13_size0.31, 1) != slice(0, $4load_attr.1.21_size0.39, 1)
----------------------------- Before Optimisation ------------------------------
--------------------------------------------------------------------------------
------------------------------ After Optimisation ------------------------------
Parallel structure is already optimal.
--------------------------------------------------------------------------------
--------------------------------------------------------------------------------

---------------------------Loop invariant code motion---------------------------
Allocation hoisting:
No allocation hoisting found

Instruction hoisting:
loop #0:
  Failed to hoist the following:
    dependency: $arg_out_var.51 = getitem(value=_4load__attr_1_13, index=$parfor__index_46.149)
    dependency: $arg_out_var.52 = getitem(value=_8load__attr_3_15, index=$parfor__index_46.149)
    dependency: $expr_out_var.50 = $arg_out_var.51 // $arg_out_var.52
loop #1:
  Failed to hoist the following:
    dependency: $arg_out_var.58 = getitem(value=_4load__attr_1_21, index=$parfor__index_53.187)
    dependency: $arg_out_var.59 = getitem(value=_8load__attr_3_23, index=$parfor__index_53.187)
    dependency: $expr_out_var.57 = $arg_out_var.58 ** $arg_out_var.59
--------------------------------------------------------------------------------

In [3]: test_impl2.parallel_diagnostics(level=4)

================================================================================
 Parallel Accelerator Optimizing:  Function test_impl2,
/Users/vhaenel/git/numba/foo031.py (106)
================================================================================


Parallel loop listing for  Function test_impl2, /Users/vhaenel/git/numba/foo031.py (106)
----------------------------------------------------------------------------|loop #ID
@numba.njit(parallel=True)                                                  |
def test_impl2(data, other_data):                                           |
    # Parallel for-loop #3 is fused into for-loop #2.                       |
    # After fusion, function test_impl2 has 1 parallel for-loop(s) #{2}.    |
    res1 = data // other_data-----------------------------------------------| #2
    res2 = data ** other_data-----------------------------------------------| #3
                                                                            |
    return res1, res2                                                       |
--------------------------------- Fusing loops ---------------------------------
Attempting fusion of parallel loops (combines loops with similar properties)...
  Trying to fuse loops #2 and #3:
    - fusion succeeded: parallel for-loop #3 is fused into for-loop #2.
----------------------------- Before Optimisation ------------------------------
Parallel region 0:
+--2 (parallel)
+--3 (parallel)


--------------------------------------------------------------------------------
------------------------------ After Optimisation ------------------------------
Parallel region 0:
+--2 (parallel, fused with loop(s): 3)



Parallel region 0 (loop #2) had 1 loop(s) fused.
--------------------------------------------------------------------------------
--------------------------------------------------------------------------------

---------------------------Loop invariant code motion---------------------------
Allocation hoisting:
No allocation hoisting found

Instruction hoisting:
loop #2:
  Failed to hoist the following:
    dependency: $arg_out_var.210 = getitem(value=data, index=$parfor__index_205.293)
    dependency: $arg_out_var.211 = getitem(value=other__data, index=$parfor__index_205.293)
    dependency: $expr_out_var.209 = $arg_out_var.210 // $arg_out_var.211
    dependency: $arg_out_var.217 = getitem(value=data, index=$parfor__index_205.293)
    dependency: $arg_out_var.218 = getitem(value=other__data, index=$parfor__index_205.293)
    dependency: $expr_out_var.216 = $arg_out_var.217 ** $arg_out_var.218
--------------------------------------------------------------------------------
```
Seems like this might be a hint:

```
Attempting fusion of parallel loops (combines loops with similar properties)...
  Trying to fuse loops #0 and #1:
    - fusion failed: loop dimension mismatched in axis 0. slice(0,
$4load_attr.1.13_size0.31, 1) != slice(0, $4load_attr.1.21_size0.39, 1)
```
@DrTodd13 could you perhaps have a look at this one, thanks!
@densmirn I have a fix for this issue in the issue5499 branch of IntelLabs/Numba.
@esc @stuartarchibald How do we get this fix integrated with master?
@DrTodd13 not sure, perhaps you could cherry-pick the fix into a fresh branch and submit a pull-request?

Resolving patch:
diff --git a/numba/core/ir_utils.py b/numba/core/ir_utils.py
--- a/numba/core/ir_utils.py
+++ b/numba/core/ir_utils.py
@@ -737,6 +737,13 @@ def is_const_call(module_name, func_name):
 alias_analysis_extensions = {}
 alias_func_extensions = {}
 
+def get_canonical_alias(v, alias_map):
+    if v not in alias_map:
+        return v
+
+    v_aliases = sorted(list(alias_map[v]))
+    return v_aliases[0]
+
 def find_potential_aliases(blocks, args, typemap, func_ir, alias_map=None,
                                                            arg_aliases=None):
     "find all array aliases and argument aliases to avoid remove as dead"
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
@@ -19,14 +19,17 @@
     find_const,
     is_namedtuple_class,
     build_definitions,
+    find_potential_aliases,
+    get_canonical_alias,
     GuardException,
 )
 from numba.core.analysis import compute_cfg_from_blocks
 from numba.core.typing import npydecl, signature
 import copy
-from numba.core.extending import intrinsic
+from numba.core.extending import intrinsic, register_jitable
 import llvmlite.llvmpy.core as lc
 import llvmlite
+from numba.np.unsafe.ndarray import to_fixed_tuple
 
 UNKNOWN_CLASS = -1
 CONST_CLASS = 0
@@ -138,10 +141,12 @@ def assert_equiv(typingctx, *val):
     assert len(val[0]) > 1
     # Arguments must be either array, tuple, or integer
     assert all(
-        isinstance(a, types.ArrayCompatible)
-        or isinstance(a, types.BaseTuple)
-        or isinstance(a, types.SliceType)
-        or isinstance(a, types.Integer)
+        isinstance(a, (
+            types.ArrayCompatible,
+            types.BaseTuple,
+            types.SliceType,
+            types.Integer
+        ))
         for a in val[0][1:]
     )
     if not isinstance(val[0][0], types.StringLiteral):
@@ -606,9 +611,9 @@ def get_shape_classes(self, name):
             name = name.name
         typ = self.typemap[name] if name in self.typemap else None
         if not (
-            isinstance(typ, types.BaseTuple)
-            or isinstance(typ, types.SliceType)
-            or isinstance(typ, types.ArrayCompatible)
+            isinstance(typ, (
+                types.BaseTuple, types.SliceType, types.ArrayCompatible
+            ))
         ):
             return []
         # Treat 0d arrays like scalars.
@@ -1046,6 +1051,27 @@ 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
 
@@ -1066,6 +1092,8 @@ def __init__(self, context, func_ir, typemap, calltypes):
         self.equiv_sets = {}
         # keep attr calls to arrays like t=A.sum() as {t:('sum',A)}
         self.array_attr_calls = {}
+        # keep attrs of objects (value,attr)->shape_var
+        self.object_attrs = {}
         # keep prepended instructions from conditional branch
         self.prepends = {}
         # keep track of pruned precessors when branch degenerates to jump
@@ -1102,6 +1130,13 @@ def run(self, blocks=None, equiv_set=None):
         else:
             init_equiv_set = equiv_set
 
+        self.alias_map, self.arg_aliases = find_potential_aliases(
+            blocks,
+            self.func_ir.arg_names,
+            self.typemap,
+            self.func_ir
+        )
+
         aa_count_save = ArrayAnalysis.aa_count
         ArrayAnalysis.aa_count += 1
         if config.DEBUG_ARRAY_OPT >= 1:
@@ -1238,8 +1273,11 @@ def _define(self, equiv_set, var, typ, value):
         self.typemap[var.name] = typ
         self.func_ir._definitions[var.name] = [value]
         redefineds = set()
-        # Propagate the success or failure of define.
-        return equiv_set.define(var, redefineds, self.func_ir, typ)
+        equiv_set.define(var, redefineds, self.func_ir, typ)
+
+    class AnalyzeResult(object):
+        def __init__(self, **kwargs):
+            self.kwargs = kwargs
 
     def _analyze_inst(self, label, scope, equiv_set, inst, redefined):
         pre = []
@@ -1253,17 +1291,25 @@ def _analyze_inst(self, label, scope, equiv_set, inst, redefined):
             if isinstance(typ, types.ArrayCompatible) and typ.ndim == 0:
                 shape = ()
             elif isinstance(inst.value, ir.Expr):
-                result = self._analyze_expr(scope, equiv_set, inst.value)
+                result = self._analyze_expr(scope, equiv_set, inst.value, lhs)
                 if result:
-                    shape = result[0]
-                    pre = result[1]
-                    if len(result) > 2:
-                        rhs = result[2]
-                        inst.value = rhs
-            elif isinstance(inst.value, ir.Var) or isinstance(
-                inst.value, ir.Const
-            ):
+                    require(isinstance(result, ArrayAnalysis.AnalyzeResult))
+                    if 'shape' in result.kwargs:
+                        shape = result.kwargs['shape']
+                    if 'pre' in result.kwargs:
+                        pre.extend(result.kwargs['pre'])
+                    if 'post' in result.kwargs:
+                        post.extend(result.kwargs['post'])
+                    if 'rhs' in result.kwargs:
+                        inst.value = result.kwargs['rhs']
+            elif isinstance(inst.value, (ir.Var, ir.Const)):
                 shape = inst.value
+            elif isinstance(inst.value, ir.Global):
+                gvalue = inst.value.value
+                if isinstance(gvalue, tuple):
+                    shape = gvalue
+                elif isinstance(gvalue, int):
+                    shape = (gvalue,)
 
             if isinstance(shape, ir.Const):
                 if isinstance(shape.value, tuple):
@@ -1310,13 +1356,13 @@ def _analyze_inst(self, label, scope, equiv_set, inst, redefined):
                         and not equiv_set.has_shape(shape)
                     )
                 ):
-                    (shape, post) = self._gen_shape_call(
-                        equiv_set, lhs, typ.ndim, shape
+                    shape = self._gen_shape_call(
+                        equiv_set, lhs, typ.ndim, shape, post
                     )
             elif isinstance(typ, types.UniTuple):
                 if shape and isinstance(typ.dtype, types.Integer):
-                    (shape, post) = self._gen_shape_call(
-                        equiv_set, lhs, len(typ), shape
+                    shape = self._gen_shape_call(
+                        equiv_set, lhs, len(typ), shape, post
                     )
 
             """ See the comment on the define() function.
@@ -1338,9 +1384,7 @@ def _analyze_inst(self, label, scope, equiv_set, inst, redefined):
                 needs_define = equiv_set.insert_equiv(lhs, shape)
             if needs_define:
                 equiv_set.define(lhs, redefined, self.func_ir, typ)
-        elif isinstance(inst, ir.StaticSetItem) or isinstance(
-            inst, ir.SetItem
-        ):
+        elif isinstance(inst, (ir.StaticSetItem, ir.SetItem)):
             index = (
                 inst.index if isinstance(inst, ir.SetItem) else inst.index_var
             )
@@ -1353,7 +1397,9 @@ def _analyze_inst(self, label, scope, equiv_set, inst, redefined):
                 assert isinstance(inst, ir.SetItem)
                 inst.index = result[0]
             result = result[1]
-            (target_shape, pre) = result
+            target_shape = result.kwargs['shape']
+            if 'pre' in result.kwargs:
+                pre = result.kwargs['pre']
             value_shape = equiv_set.get_shape(inst.value)
             if value_shape == ():  # constant
                 equiv_set.set_shape_setitem(inst, target_shape)
@@ -1364,9 +1410,13 @@ def _analyze_inst(self, label, scope, equiv_set, inst, redefined):
                 target_ndim = target_typ.ndim
                 shapes = [target_shape, value_shape]
                 names = [inst.target.name, inst.value.name]
-                shape, asserts = self._broadcast_assert_shapes(
+                broadcast_result = self._broadcast_assert_shapes(
                     scope, equiv_set, inst.loc, shapes, names
                 )
+                require('shape' in broadcast_result.kwargs)
+                require('pre' in broadcast_result.kwargs)
+                shape = broadcast_result.kwargs['shape']
+                asserts = broadcast_result.kwargs['pre']
                 n = len(shape)
                 # shape dimension must be within target dimension
                 assert target_ndim >= n
@@ -1467,15 +1517,15 @@ def handle_call_binop(cond_def):
 
         return pre, post
 
-    def _analyze_expr(self, scope, equiv_set, expr):
+    def _analyze_expr(self, scope, equiv_set, expr, lhs):
         fname = "_analyze_op_{}".format(expr.op)
         try:
             fn = getattr(self, fname)
         except AttributeError:
             return None
-        return guard(fn, scope, equiv_set, expr)
+        return guard(fn, scope, equiv_set, expr, lhs)
 
-    def _analyze_op_getattr(self, scope, equiv_set, expr):
+    def _analyze_op_getattr(self, scope, equiv_set, expr, lhs):
         # TODO: getattr of npytypes.Record
         if expr.attr == "T" and self._isarray(expr.value.name):
             return self._analyze_op_call_numpy_transpose(
@@ -1483,19 +1533,36 @@ def _analyze_op_getattr(self, scope, equiv_set, expr):
             )
         elif expr.attr == "shape":
             shape = equiv_set.get_shape(expr.value)
-            return shape, []
+            return ArrayAnalysis.AnalyzeResult(shape=shape)
+        elif self._isarray(lhs.name):
+            canonical_value = get_canonical_alias(
+                expr.value.name, self.alias_map
+            )
+            if (canonical_value, expr.attr) in self.object_attrs:
+                return ArrayAnalysis.AnalyzeResult(
+                    shape=self.object_attrs[(canonical_value, expr.attr)]
+                )
+            else:
+                typ = self.typemap[lhs.name]
+                post = []
+                shape = self._gen_shape_call(
+                    equiv_set, lhs, typ.ndim, None, post
+                )
+                self.object_attrs[(canonical_value, expr.attr)] = shape
+                return ArrayAnalysis.AnalyzeResult(shape=shape, post=post)
+
         return None
 
-    def _analyze_op_cast(self, scope, equiv_set, expr):
-        return expr.value, []
+    def _analyze_op_cast(self, scope, equiv_set, expr, lhs):
+        return ArrayAnalysis.AnalyzeResult(shape=expr.value)
 
-    def _analyze_op_exhaust_iter(self, scope, equiv_set, expr):
+    def _analyze_op_exhaust_iter(self, scope, equiv_set, expr, lhs):
         var = expr.value
         typ = self.typemap[var.name]
         if isinstance(typ, types.BaseTuple):
             require(len(typ) == expr.count)
             require(equiv_set.has_shape(var))
-            return var, []
+            return ArrayAnalysis.AnalyzeResult(shape=var)
         return None
 
     def gen_literal_slice_part(
@@ -2012,15 +2079,16 @@ def to_shape(typ, index, dsize):
         shape = tuple(shape_list)
         require(not all(x is None for x in shape))
         shape = tuple(x for x in shape if x is not None)
-        return (replacement_build_tuple_var, (shape, stmts))
+        return (replacement_build_tuple_var,
+                ArrayAnalysis.AnalyzeResult(shape=shape, pre=stmts))
 
-    def _analyze_op_getitem(self, scope, equiv_set, expr):
+    def _analyze_op_getitem(self, scope, equiv_set, expr, lhs):
         result = self._index_to_shape(scope, equiv_set, expr.value, expr.index)
         if result[0] is not None:
             expr.index = result[0]
         return result[1]
 
-    def _analyze_op_static_getitem(self, scope, equiv_set, expr):
+    def _analyze_op_static_getitem(self, scope, equiv_set, expr, lhs):
         var = expr.value
         typ = self.typemap[var.name]
         if not isinstance(typ, types.BaseTuple):
@@ -2033,37 +2101,37 @@ def _analyze_op_static_getitem(self, scope, equiv_set, expr):
         shape = equiv_set._get_shape(var)
         if isinstance(expr.index, int):
             require(expr.index < len(shape))
-            return shape[expr.index], []
+            return ArrayAnalysis.AnalyzeResult(shape=shape[expr.index])
         elif isinstance(expr.index, slice):
-            return shape[expr.index], []
+            return ArrayAnalysis.AnalyzeResult(shape=shape[expr.index])
         require(False)
 
-    def _analyze_op_unary(self, scope, equiv_set, expr):
+    def _analyze_op_unary(self, scope, equiv_set, expr, lhs):
         require(expr.fn in UNARY_MAP_OP)
         # for scalars, only + operator results in equivalence
         # for example, if "m = -n", m and n are not equivalent
         if self._isarray(expr.value.name) or expr.fn == operator.add:
-            return expr.value, []
+            return ArrayAnalysis.AnalyzeResult(shape=expr.value)
         return None
 
-    def _analyze_op_binop(self, scope, equiv_set, expr):
+    def _analyze_op_binop(self, scope, equiv_set, expr, lhs):
         require(expr.fn in BINARY_MAP_OP)
         return self._analyze_broadcast(
             scope, equiv_set, expr.loc, [expr.lhs, expr.rhs], expr.fn
         )
 
-    def _analyze_op_inplace_binop(self, scope, equiv_set, expr):
+    def _analyze_op_inplace_binop(self, scope, equiv_set, expr, lhs):
         require(expr.fn in INPLACE_BINARY_MAP_OP)
         return self._analyze_broadcast(
             scope, equiv_set, expr.loc, [expr.lhs, expr.rhs], expr.fn
         )
 
-    def _analyze_op_arrayexpr(self, scope, equiv_set, expr):
+    def _analyze_op_arrayexpr(self, scope, equiv_set, expr, lhs):
         return self._analyze_broadcast(
             scope, equiv_set, expr.loc, expr.list_vars(), None
         )
 
-    def _analyze_op_build_tuple(self, scope, equiv_set, expr):
+    def _analyze_op_build_tuple(self, scope, equiv_set, expr, lhs):
         # For the moment, we can't do anything with tuples that
         # contain multi-dimensional arrays, compared to array dimensions.
         # Return None to say we won't track this tuple if a part of it
@@ -2085,11 +2153,14 @@ def _analyze_op_build_tuple(self, scope, equiv_set, expr):
                 break
         else:
             out = tuple([ir.Const(x, expr.loc) for x in consts])
-            return out, [], ir.Const(tuple(consts), expr.loc)
+            return ArrayAnalysis.AnalyzeResult(
+                shape=out,
+                rhs=ir.Const(tuple(consts), expr.loc)
+            )
         # default return for non-const
-        return tuple(expr.items), []
+        return ArrayAnalysis.AnalyzeResult(shape=tuple(expr.items))
 
-    def _analyze_op_call(self, scope, equiv_set, expr):
+    def _analyze_op_call(self, scope, equiv_set, expr, lhs):
         from numba.stencils.stencil import StencilFunc
 
         callee = expr.func
@@ -2097,7 +2168,7 @@ def _analyze_op_call(self, scope, equiv_set, expr):
         if isinstance(
             callee_def, (ir.Global, ir.FreeVar)
         ) and is_namedtuple_class(callee_def.value):
-            return tuple(expr.args), []
+            return ArrayAnalysis.AnalyzeResult(shape=tuple(expr.args))
         if isinstance(callee_def, (ir.Global, ir.FreeVar)) and isinstance(
             callee_def.value, StencilFunc
         ):
@@ -2159,7 +2230,7 @@ def _analyze_op_call_builtins_len(self, scope, equiv_set, loc, args, kws):
         typ = self.typemap[var.name]
         require(isinstance(typ, types.ArrayCompatible))
         shape = equiv_set._get_shape(var)
-        return shape[0], [], shape[0]
+        return ArrayAnalysis.AnalyzeResult(shape=shape[0], rhs=shape[0])
 
     def _analyze_op_call_numba_parfors_array_analysis_assert_equiv(
         self, scope, equiv_set, loc, args, kws
@@ -2189,12 +2260,14 @@ def _analyze_op_call_numba_parfors_array_analysis_wrap_index(
             vs = equiv_set.ind_to_var[wrap_ind]
             require(vs != [])
             # Return the shape of the variable from the previous wrap_index.
-            return ((vs[0],), [])
+            return ArrayAnalysis.AnalyzeResult(shape=(vs[0],))
         else:
             # We haven't seen this combination of slice and dim
             # equivalence class ids so return a WrapIndexMeta so that
             # _analyze_inst can establish the connection to the lhs var.
-            return (WrapIndexMeta(slice_eq, dim_eq), [])
+            return ArrayAnalysis.AnalyzeResult(
+                shape=WrapIndexMeta(slice_eq, dim_eq)
+            )
 
     def _analyze_numpy_create_array(self, scope, equiv_set, loc, args, kws):
         shape_var = None
@@ -2203,7 +2276,7 @@ def _analyze_numpy_create_array(self, scope, equiv_set, loc, args, kws):
         elif "shape" in kws:
             shape_var = kws["shape"]
         if shape_var:
-            return shape_var, []
+            return ArrayAnalysis.AnalyzeResult(shape=shape_var)
         raise errors.UnsupportedRewriteError(
             "Must specify a shape for array creation",
             loc=loc,
@@ -2245,14 +2318,14 @@ def _analyze_op_call_numpy_eye(self, scope, equiv_set, loc, args, kws):
             M = kws["M"]
         else:
             M = N
-        return (N, M), []
+        return ArrayAnalysis.AnalyzeResult(shape=(N, M))
 
     def _analyze_op_call_numpy_identity(
         self, scope, equiv_set, loc, args, kws
     ):
         assert len(args) > 0
         N = args[0]
-        return (N, N), []
+        return ArrayAnalysis.AnalyzeResult(shape=(N, N))
 
     def _analyze_op_call_numpy_diag(self, scope, equiv_set, loc, args, kws):
         # We can only reason about the output shape when the input is 1D or
@@ -2269,10 +2342,10 @@ def _analyze_op_call_numpy_diag(self, scope, equiv_set, loc, args, kws):
                         return None
                 (m, n) = equiv_set._get_shape(a)
                 if equiv_set.is_equiv(m, n):
-                    return (m,), []
+                    return ArrayAnalysis.AnalyzeResult(shape=(m,))
             elif atyp.ndim == 1:
                 (m,) = equiv_set._get_shape(a)
-                return (m, m), []
+                return ArrayAnalysis.AnalyzeResult(shape=(m, m))
         return None
 
     def _analyze_numpy_array_like(self, scope, equiv_set, args, kws):
@@ -2280,11 +2353,11 @@ def _analyze_numpy_array_like(self, scope, equiv_set, args, kws):
         var = args[0]
         typ = self.typemap[var.name]
         if isinstance(typ, types.Integer):
-            return (1,), []
+            return ArrayAnalysis.AnalyzeResult(shape=(1,))
         elif isinstance(typ, types.ArrayCompatible) and equiv_set.has_shape(
             var
         ):
-            return var, []
+            return ArrayAnalysis.AnalyzeResult(shape=var)
         return None
 
     def _analyze_op_call_numpy_ravel(self, scope, equiv_set, loc, args, kws):
@@ -2297,9 +2370,9 @@ def _analyze_op_call_numpy_ravel(self, scope, equiv_set, loc, args, kws):
             if typ.layout == "C":
                 # output is the same as input (no copy) for 'C' layout
                 # optimize out the call
-                return var, [], var
+                return ArrayAnalysis.AnalyzeResult(shape=var, rhs=var)
             else:
-                return var, []
+                return ArrayAnalysis.AnalyzeResult(shape=var)
         # TODO: handle multi-D input arrays (calc array size)
         return None
 
@@ -2337,7 +2410,7 @@ def _analyze_op_call_numpy_reshape(self, scope, equiv_set, loc, args, kws):
         if n == 2:
             typ = self.typemap[args[1].name]
             if isinstance(typ, types.BaseTuple):
-                return args[1], []
+                return ArrayAnalysis.AnalyzeResult(shape=args[1])
 
         # Reshape is allowed to take one argument that has the value <0.
         # This means that the size of that dimension should be inferred from
@@ -2403,7 +2476,7 @@ def _analyze_op_call_numpy_reshape(self, scope, equiv_set, loc, args, kws):
             # replacing the negative.
             args[neg_one_index] = calc_size_var
 
-        return tuple(args[1:]), stmts
+        return ArrayAnalysis.AnalyzeResult(shape=tuple(args[1:]), pre=stmts)
 
     def _analyze_op_call_numpy_transpose(
         self, scope, equiv_set, loc, args, kws
@@ -2415,20 +2488,20 @@ def _analyze_op_call_numpy_transpose(
         ), "Invalid np.transpose argument"
         shape = equiv_set._get_shape(in_arr)
         if len(args) == 1:
-            return tuple(reversed(shape)), []
+            return ArrayAnalysis.AnalyzeResult(shape=tuple(reversed(shape)))
         axes = [guard(find_const, self.func_ir, a) for a in args[1:]]
         if isinstance(axes[0], tuple):
             axes = list(axes[0])
         if None in axes:
             return None
         ret = [shape[i] for i in axes]
-        return tuple(ret), []
+        return ArrayAnalysis.AnalyzeResult(shape=tuple(ret))
 
     def _analyze_op_call_numpy_random_rand(
         self, scope, equiv_set, loc, args, kws
     ):
         if len(args) > 0:
-            return tuple(args), []
+            return ArrayAnalysis.AnalyzeResult(shape=tuple(args))
         return None
 
     def _analyze_op_call_numpy_random_randn(
@@ -2442,9 +2515,9 @@ def _analyze_op_numpy_random_with_size(
         self, pos, scope, equiv_set, args, kws
     ):
         if "size" in kws:
-            return kws["size"], []
+            return ArrayAnalysis.AnalyzeResult(shape=kws["size"])
         if len(args) > pos:
-            return args[pos], []
+            return ArrayAnalysis.AnalyzeResult(shape=args[pos])
         return None
 
     def _analyze_op_call_numpy_random_ranf(
@@ -2650,7 +2723,10 @@ def _analyze_op_call_numpy_concatenate(
                     )
                     size = sizes[0]
                 new_shape.append(size)
-        return tuple(new_shape), sum(asserts, [])
+        return ArrayAnalysis.AnalyzeResult(
+            shape=tuple(new_shape),
+            pre=sum(asserts, [])
+        )
 
     def _analyze_op_call_numpy_stack(self, scope, equiv_set, loc, args, kws):
         assert len(args) > 0
@@ -2676,7 +2752,7 @@ def _analyze_op_call_numpy_stack(self, scope, equiv_set, loc, args, kws):
             axis = len(shape) + axis + 1
         require(0 <= axis <= len(shape))
         new_shape = list(shape[0:axis]) + [n] + list(shape[axis:])
-        return tuple(new_shape), asserts
+        return ArrayAnalysis.AnalyzeResult(shape=tuple(new_shape), pre=asserts)
 
     def _analyze_op_call_numpy_vstack(self, scope, equiv_set, loc, args, kws):
         assert len(args) == 1
@@ -2723,9 +2799,8 @@ def _analyze_op_call_numpy_dstack(self, scope, equiv_set, loc, args, kws):
                 scope, equiv_set, loc, args, kws
             )
             require(result)
-            (shape, pre) = result
-            shape = tuple([1] + list(shape))
-            return shape, pre
+            result.kwargs['shape'] = tuple([1] + list(result.kwargs['shape']))
+            return result
         elif typ.ndim == 2:
             kws["axis"] = 2
             return self._analyze_op_call_numpy_stack(
@@ -2754,7 +2829,7 @@ def _analyze_op_call_numpy_linspace(
             num = args[2]
         elif "num" in kws:
             num = kws["num"]
-        return (num,), []
+        return ArrayAnalysis.AnalyzeResult(shape=(num,))
 
     def _analyze_op_call_numpy_dot(self, scope, equiv_set, loc, args, kws):
         n = len(args)
@@ -2771,17 +2846,26 @@ def _analyze_op_call_numpy_dot(self, scope, equiv_set, loc, args, kws):
             asserts = self._call_assert_equiv(
                 scope, loc, equiv_set, [shapes[0][0], shapes[1][-2]]
             )
-            return tuple(shapes[1][0:-2] + shapes[1][-1:]), asserts
+            return ArrayAnalysis.AnalyzeResult(
+                shape=tuple(shapes[1][0:-2] + shapes[1][-1:]),
+                pre=asserts
+            )
         if dims[1] == 1:
             asserts = self._call_assert_equiv(
                 scope, loc, equiv_set, [shapes[0][-1], shapes[1][0]]
             )
-            return tuple(shapes[0][0:-1]), asserts
+            return ArrayAnalysis.AnalyzeResult(
+                shape=tuple(shapes[0][0:-1]),
+                pre=asserts
+            )
         if dims[0] == 2 and dims[1] == 2:
             asserts = self._call_assert_equiv(
                 scope, loc, equiv_set, [shapes[0][1], shapes[1][0]]
             )
-            return (shapes[0][0], shapes[1][1]), asserts
+            return ArrayAnalysis.AnalyzeResult(
+                shape=(shapes[0][0], shapes[1][1]),
+                pre=asserts
+            )
         if dims[0] > 2:  # TODO: handle higher dimension cases
             pass
         return None
@@ -2805,13 +2889,65 @@ def _analyze_stencil(self, scope, equiv_set, stencil_func, loc, args, kws):
         require(n > 0)
         asserts = self._call_assert_equiv(scope, loc, equiv_set, rel_idx_arrs)
         shape = equiv_set.get_shape(rel_idx_arrs[0])
-        return shape, asserts
+        return ArrayAnalysis.AnalyzeResult(shape=shape, pre=asserts)
 
     def _analyze_op_call_numpy_linalg_inv(
         self, scope, equiv_set, loc, args, kws
     ):
         require(len(args) >= 1)
-        return equiv_set._get_shape(args[0]), []
+        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
@@ -2826,16 +2962,22 @@ def _analyze_broadcast(self, scope, equiv_set, loc, args, fn):
             tup0typ = self.typemap[tups[0].name]
             tup1typ = self.typemap[tups[1].name]
             if tup0typ.count == 0:
-                return (equiv_set.get_shape(tups[1]), [])
+                return ArrayAnalysis.AnalyzeResult(
+                    shape=equiv_set.get_shape(tups[1])
+                )
             if tup1typ.count == 0:
-                return (equiv_set.get_shape(tups[0]), [])
+                return ArrayAnalysis.AnalyzeResult(
+                    shape=equiv_set.get_shape(tups[0])
+                )
 
             try:
                 shapes = [equiv_set.get_shape(x) for x in tups]
                 if None in shapes:
                     return None
                 concat_shapes = sum(shapes, ())
-                return (concat_shapes, [])
+                return ArrayAnalysis.AnalyzeResult(
+                    shape=concat_shapes
+                )
             except GuardException:
                 return None
 
@@ -2849,13 +2991,18 @@ def _analyze_broadcast(self, scope, equiv_set, loc, args, fn):
         try:
             shapes = [equiv_set.get_shape(x) for x in arrs]
         except GuardException:
-            return (
-                arrs[0],
-                self._call_assert_equiv(scope, loc, equiv_set, arrs),
+            return ArrayAnalysis.AnalyzeResult(
+                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
             )
-        return self._broadcast_assert_shapes(
-            scope, equiv_set, loc, shapes, names
-        )
 
     def _broadcast_assert_shapes(self, scope, equiv_set, loc, shapes, names):
         """Produce assert_equiv for sizes in each dimension, taking into
@@ -2887,7 +3034,10 @@ def _broadcast_assert_shapes(self, scope, equiv_set, loc, shapes, names):
                 )
             )
             new_shape.append(sizes[0])
-        return tuple(reversed(new_shape)), sum(asserts, [])
+        return ArrayAnalysis.AnalyzeResult(
+            shape=tuple(reversed(new_shape)),
+            pre=sum(asserts, [])
+        )
 
     def _call_assert_equiv(self, scope, loc, equiv_set, args, names=None):
         insts = self._make_assert_equiv(
@@ -2959,8 +3109,7 @@ def _make_assert_equiv(self, scope, loc, equiv_set, _args, names=None):
             ir.Assign(value=value, target=var, loc=loc),
         ]
 
-    def _gen_shape_call(self, equiv_set, var, ndims, shape):
-        out = []
+    def _gen_shape_call(self, equiv_set, var, ndims, shape, post):
         # attr call: A_sh_attr = getattr(A, shape)
         if isinstance(shape, ir.Var):
             shape = equiv_set.get_shape(shape)
@@ -2987,9 +3136,7 @@ def _gen_shape_call(self, equiv_set, var, ndims, shape):
             if shape and shape[i]:
                 if isinstance(shape[i], ir.Var):
                     typ = self.typemap[shape[i].name]
-                    if isinstance(typ, types.Number) or isinstance(
-                        typ, types.SliceType
-                    ):
+                    if isinstance(typ, (types.Number, types.SliceType)):
                         size_var = shape[i]
                         skip = True
                 else:
@@ -3003,7 +3150,7 @@ def _gen_shape_call(self, equiv_set, var, ndims, shape):
                         mk_unique_var("{}_size{}".format(var.name, i)),
                         var.loc,
                     )
-                    out.append(ir.Assign(size_val, size_var, var.loc))
+                    post.append(ir.Assign(size_val, size_var, var.loc))
                     self._define(equiv_set, size_var, types.intp, size_val)
                     skip = True
             if not skip:
@@ -3016,14 +3163,14 @@ def _gen_shape_call(self, equiv_set, var, ndims, shape):
                 getitem = ir.Expr.static_getitem(attr_var, i, None, var.loc)
                 use_attr_var = True
                 self.calltypes[getitem] = None
-                out.append(ir.Assign(getitem, size_var, var.loc))
+                post.append(ir.Assign(getitem, size_var, var.loc))
                 self._define(equiv_set, size_var, types.intp, getitem)
             size_vars.append(size_var)
         if use_attr_var and shape_attr_call:
             # only insert shape call if there is any getitem call
-            out.insert(0, ir.Assign(shape_attr_call, attr_var, var.loc))
+            post.insert(0, ir.Assign(shape_attr_call, attr_var, var.loc))
             self._define(equiv_set, attr_var, shape_attr_typ, shape_attr_call)
-        return tuple(size_vars), out
+        return tuple(size_vars)
 
     def _isarray(self, varname):
         typ = self.typemap[varname]


Test patch: