# Buggy code
def replace_return_with_setitem(self, blocks, index_vars, out_name):
    """
    Find return statements in the IR and replace them with a SetItem
    call of the value "returned" by the kernel into the result array.
    """
    for block in blocks.values():
        scope = block.scope
        loc = block.loc
        new_body = []
        for stmt in block.body:
            if isinstance(stmt, ir.Return):
                # If 1D array then avoid the tuple construction.
                if len(index_vars) == 1:
                    rvar = ir.Var(scope, out_name, loc)
                    ivar = ir.Var(scope, index_vars[0], loc)
                    new_body.append(ir.SetItem(rvar, ivar, stmt.value, loc))
                else:
                    # Convert the string names of the index variables into
                    # ir.Var's.
                    var_index_vars = []
                    for one_var in index_vars:
                        index_var = ir.Var(scope, one_var, loc)
                        var_index_vars += [index_var]

                    s_index_name = ir_utils.mk_unique_var("stencil_index")
                    s_index_var = ir.Var(scope, s_index_name, loc)
                    # Build a tuple from the index ir.Var's.
                    tuple_call = ir.Expr.build_tuple(var_index_vars, loc)
                    new_body.append(ir.Assign(tuple_call, s_index_var, loc))
                    rvar = ir.Var(scope, out_name, loc)
                    # Write the return statements original value into
                    # the array using the tuple index.
                    si = ir.SetItem(rvar, s_index_var, stmt.value, loc)
                    new_body.append(si)
            else:
                new_body.append(stmt)
        block.body = new_body


# Traceback (AttributeError)
---------------------------------------------------------------------------
AttributeError                            Traceback (most recent call last)
<ipython-input-24-95489a051be1> in <module>()
----> 1 kernel1(np.array([[1,2,3], [4,5,6], [7,8,9]]))

~/envs/py36/lib/python3.6/site-packages/numba/stencil.py in __call__(self, *args, **kwargs)
682         (real_ret, typemap, calltypes) = self.get_return_type(array_types)
683         new_func = self._stencil_wrapper(result, None, real_ret, typemap,
--> 684                                          calltypes, *array_types_full)
685
686         if result is None:

~/envs/py36/lib/python3.6/site-packages/numba/stencil.py in _stencil_wrapper(self, result, sigret, return_type, typemap, calltypes, *args)
632             break
633
--> 634         stencil_ir.blocks = ir_utils.rename_labels(stencil_ir.blocks)
635         ir_utils.remove_dels(stencil_ir.blocks)
636

~/envs/py36/lib/python3.6/site-packages/numba/ir_utils.py in rename_labels(blocks)
1032     lowering requires this order.
1033     """
-> 1034     topo_order = find_topo_order(blocks)
1035
1036     # make a block with return last if available (just for readability)

~/envs/py36/lib/python3.6/site-packages/numba/ir_utils.py in find_topo_order(blocks, cfg)
925     """
926     if cfg == None:
--> 927         cfg = compute_cfg_from_blocks(blocks)
928     post_order = []
929     seen = set()

~/envs/py36/lib/python3.6/site-packages/numba/analysis.py in compute_cfg_from_blocks(blocks)
210     for k, b in blocks.items():
211         term = b.terminator
--> 212         for target in term.get_targets():
213             cfg.add_edge(k, target)
214

# Fixed code
def replace_return_with_setitem(self, blocks, index_vars, out_name):
    """
    Find return statements in the IR and replace them with a SetItem
    call of the value "returned" by the kernel into the result array.
    Returns the block labels that contained return statements.
    """
    ret_blocks = []

    for label, block in blocks.items():
        scope = block.scope
        loc = block.loc
        new_body = []
        for stmt in block.body:
            if isinstance(stmt, ir.Return):
                ret_blocks.append(label)
                # If 1D array then avoid the tuple construction.
                if len(index_vars) == 1:
                    rvar = ir.Var(scope, out_name, loc)
                    ivar = ir.Var(scope, index_vars[0], loc)
                    new_body.append(ir.SetItem(rvar, ivar, stmt.value, loc))
                else:
                    # Convert the string names of the index variables into
                    # ir.Var's.
                    var_index_vars = []
                    for one_var in index_vars:
                        index_var = ir.Var(scope, one_var, loc)
                        var_index_vars += [index_var]

                    s_index_name = ir_utils.mk_unique_var("stencil_index")
                    s_index_var = ir.Var(scope, s_index_name, loc)
                    # Build a tuple from the index ir.Var's.
                    tuple_call = ir.Expr.build_tuple(var_index_vars, loc)
                    new_body.append(ir.Assign(tuple_call, s_index_var, loc))
                    rvar = ir.Var(scope, out_name, loc)
                    # Write the return statements original value into
                    # the array using the tuple index.
                    si = ir.SetItem(rvar, s_index_var, stmt.value, loc)
                    new_body.append(si)
            else:
                new_body.append(stmt)
        block.body = new_body
    return ret_blocks