stencil.py

#
# Copyright (c) 2017 Intel Corporation
# SPDX-License-Identifier: BSD-2-Clause
#

import copy

import numpy as np
from llvmlite import ir as lir

from numba import compiler, types, ir_utils, ir, typing, numpy_support, utils
from numba import config
from numba.typing.templates import (CallableTemplate, signature, infer_global,
                                    AbstractTemplate)
from numba.targets import registry
from numba.targets.imputils import lower_builtin
from numba.extending import register_jitable
from numba.six import exec_


class StencilFuncLowerer(object):
    '''Callable class responsible for lowering calls to a specific StencilFunc.
    '''
    def __init__(self, sf):
        self.stencilFunc = sf

    def __call__(self, context, builder, sig, args):
        cres = self.stencilFunc.compile_for_argtys(sig.args, {},
                    sig.return_type, None)
        return context.call_internal(builder, cres.fndesc, sig, args)

@register_jitable
def raise_if_incompatible_array_sizes(a, *args):
    ashape = a.shape
    for arg in args:
        if a.ndim != arg.ndim:
            raise ValueError("Secondary stencil array does not have same number "
                             " of dimensions as the first stencil input.")
        argshape = arg.shape
        for i in range(len(ashape)):
            if ashape[i] > argshape[i]:
                raise ValueError("Secondary stencil array has some dimension "
                                 "smaller the same dimension in the first "
                                 "stencil input.")

class StencilFunc(object):
    """
    A special type to hold stencil information for the IR.
    """

    id_counter = 0

    def __init__(self, kernel_ir, mode, options):
        self.id = type(self).id_counter
        type(self).id_counter += 1
        self.kernel_ir = kernel_ir
        self.mode = mode
        self.options = options
        self.kws = []       # remember original kws arguments

        # stencils only supported for CPU context currently
        self._typingctx = registry.cpu_target.typing_context
        self._targetctx = registry.cpu_target.target_context
        self._typingctx.refresh()
        self._targetctx.refresh()
        self._install_type(self._typingctx)
        self.neighborhood = self.options.get("neighborhood")
        self._type_cache = {}
        self._lower_me = StencilFuncLowerer(self)

    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

    def add_indices_to_kernel(self, kernel, index_names, ndim,
                              neighborhood, standard_indexed):
        """
        Transforms the stencil kernel as specified by the user into one
        that includes each dimension's index variable as part of the getitem
        calls.  So, in effect array[-1] becomes array[index0-1].
        """
        const_dict = {}
        kernel_consts = []

        if config.DEBUG_ARRAY_OPT == 1:
            print("add_indices_to_kernel", ndim, neighborhood)
            ir_utils.dump_blocks(kernel.blocks)

        if neighborhood is None:
            need_to_calc_kernel = True
        else:
            need_to_calc_kernel = False
            if len(neighborhood) != ndim:
                raise ValueError("%d dimensional neighborhood specified for %d " \
                    "dimensional input array" % (len(neighborhood), ndim))

        tuple_table = ir_utils.get_tuple_table(kernel.blocks)

        relatively_indexed = set()

        for block in kernel.blocks.values():
            scope = block.scope
            loc = block.loc
            new_body = []
            for stmt in block.body:
                if (isinstance(stmt, ir.Assign) and
                    isinstance(stmt.value, ir.Const)):
                    if config.DEBUG_ARRAY_OPT == 1:
                        print("remembering in const_dict", stmt.target.name,
                              stmt.value.value)
                    # Remember consts for use later.
                    const_dict[stmt.target.name] = stmt.value.value
                if ((isinstance(stmt, ir.Assign)
                        and isinstance(stmt.value, ir.Expr)
                        and stmt.value.op in ['setitem', 'static_setitem']
                        and stmt.value.value.name in kernel.arg_names) or
                   (isinstance(stmt, ir.SetItem)
                        and stmt.target.name in kernel.arg_names)):
                    raise ValueError("Assignments to arrays passed to stencil " \
                        "kernels is not allowed.")
                if (isinstance(stmt, ir.Assign)
                        and isinstance(stmt.value, ir.Expr)
                        and stmt.value.op in ['getitem', 'static_getitem']
                        and stmt.value.value.name in kernel.arg_names
                        and stmt.value.value.name not in standard_indexed):
                    # We found a getitem from the input array.
                    if stmt.value.op == 'getitem':
                        stmt_index_var = stmt.value.index
                    else:
                        stmt_index_var = stmt.value.index_var
                        # allow static_getitem since rewrite passes are applied
                        #raise ValueError("Unexpected static_getitem in add_indices_to_kernel.")

                    relatively_indexed.add(stmt.value.value.name)

                    # Store the index used after looking up the variable in
                    # the const dictionary.
                    if need_to_calc_kernel:
                        assert hasattr(stmt_index_var, 'name')

                        if stmt_index_var.name in tuple_table:
                            kernel_consts += [tuple_table[stmt_index_var.name]]
                        elif stmt_index_var.name in const_dict:
                            kernel_consts += [const_dict[stmt_index_var.name]]
                        else:
                            raise ValueError("stencil kernel index is not "
                                "constant, 'neighborhood' option required")

                    if ndim == 1:
                        # Single dimension always has index variable 'index0'.
                        # tmpvar will hold the real index and is computed by
                        # adding the relative offset in stmt.value.index to
                        # the current absolute location in index0.
                        index_var = ir.Var(scope, index_names[0], loc)
                        tmpname = ir_utils.mk_unique_var("stencil_index")
                        tmpvar  = ir.Var(scope, tmpname, loc)
                        acc_call = ir.Expr.binop('+', stmt_index_var,
                                                 index_var, loc)
                        new_body.append(ir.Assign(acc_call, tmpvar, loc))
                        new_body.append(ir.Assign(
                                       ir.Expr.getitem(stmt.value.value,tmpvar,loc),
                                       stmt.target,loc))
                    else:
                        index_vars = []
                        sum_results = []
                        s_index_name = ir_utils.mk_unique_var("stencil_index")
                        s_index_var  = ir.Var(scope, s_index_name, loc)
                        const_index_vars = []
                        ind_stencils = []

                        # Same idea as above but you have to extract
                        # individual elements out of the tuple indexing
                        # expression and add the corresponding index variable
                        # to them and then reconstitute as a tuple that can
                        # index the array.
                        for dim in range(ndim):
                            tmpname = ir_utils.mk_unique_var("const_index")
                            tmpvar  = ir.Var(scope, tmpname, loc)
                            new_body.append(ir.Assign(ir.Const(dim, loc),
                                                      tmpvar, loc))
                            const_index_vars += [tmpvar]
                            index_var = ir.Var(scope, index_names[dim], loc)
                            index_vars += [index_var]

                            tmpname = ir_utils.mk_unique_var("ind_stencil_index")
                            tmpvar  = ir.Var(scope, tmpname, loc)
                            ind_stencils += [tmpvar]
                            getitemname = ir_utils.mk_unique_var("getitem")
                            getitemvar  = ir.Var(scope, getitemname, loc)
                            getitemcall = ir.Expr.getitem(stmt_index_var,
                                                       const_index_vars[dim], loc)
                            new_body.append(ir.Assign(getitemcall, getitemvar, loc))
                            acc_call = ir.Expr.binop('+', getitemvar,
                                                     index_vars[dim], loc)
                            new_body.append(ir.Assign(acc_call, tmpvar, loc))

                        tuple_call = ir.Expr.build_tuple(ind_stencils, loc)
                        new_body.append(ir.Assign(tuple_call, s_index_var, loc))
                        new_body.append(ir.Assign(
                                  ir.Expr.getitem(stmt.value.value,s_index_var,loc),
                                  stmt.target,loc))
                else:
                    new_body.append(stmt)
            block.body = new_body

        if need_to_calc_kernel:
            # Find the size of the kernel by finding the maximum absolute value
            # index used in the kernel specification.
            neighborhood = [[0,0] for _ in range(ndim)]
            if len(kernel_consts) == 0:
                raise ValueError("Stencil kernel with no accesses to "
                                 "relatively indexed arrays.")

            for index in kernel_consts:
                if isinstance(index, tuple) or isinstance(index, list):
                    for i in range(len(index)):
                        te = index[i]
                        if isinstance(te, ir.Var) and te.name in const_dict:
                            te = const_dict[te.name]
                        if isinstance(te, int):
                            neighborhood[i][0] = min(neighborhood[i][0], te)
                            neighborhood[i][1] = max(neighborhood[i][1], te)
                        else:
                            raise ValueError(
                                "stencil kernel index is not constant,"
                                "'neighborhood' option required")
                    index_len = len(index)
                elif isinstance(index, int):
                    neighborhood[0][0] = min(neighborhood[0][0], index)
                    neighborhood[0][1] = max(neighborhood[0][1], index)
                    index_len = 1
                else:
                    raise ValueError(
                        "Non-tuple or non-integer used as stencil index.")
                if index_len != ndim:
                    raise ValueError(
                        "Stencil index does not match array dimensionality.")

        return (neighborhood, relatively_indexed)


    def get_return_type(self, argtys):
        if config.DEBUG_ARRAY_OPT == 1:
            print("get_return_type", argtys)
            ir_utils.dump_blocks(self.kernel_ir.blocks)

        if not isinstance(argtys[0], types.npytypes.Array):
            raise ValueError("The first argument to a stencil kernel must "
                             "be the primary input array.")

        typemap, return_type, calltypes = compiler.type_inference_stage(
                self._typingctx,
                self.kernel_ir,
                argtys,
                None,
                {})
        if isinstance(return_type, types.npytypes.Array):
            raise ValueError(
                "Stencil kernel must return a scalar and not a numpy array.")

        real_ret = types.npytypes.Array(return_type, argtys[0].ndim,
                                                     argtys[0].layout)
        return (real_ret, typemap, calltypes)

    def _install_type(self, typingctx):
        """Constructs and installs a typing class for a StencilFunc object in
        the input typing context.
        """
        _ty_cls = type('StencilFuncTyping_' +
                       str(self.id),
                       (AbstractTemplate,),
                       dict(key=self, generic=self._type_me))
        typingctx.insert_user_function(self, _ty_cls)

    def compile_for_argtys(self, argtys, kwtys, return_type, sigret):
        # look in the type cache to find if result array is passed
        (_, result, typemap, calltypes) = self._type_cache[argtys]
        new_func = self._stencil_wrapper(result, sigret, return_type,
                                         typemap, calltypes, *argtys)
        return new_func

    def _type_me(self, argtys, kwtys):
        """
        Implement AbstractTemplate.generic() for the typing class
        built by StencilFunc._install_type().
        Return the call-site signature.
        """
        if (self.neighborhood is not None and
            len(self.neighborhood) != argtys[0].ndim):
            raise ValueError("%d dimensional neighborhood specified "
                             "for %d dimensional input array" %
                             (len(self.neighborhood), argtys[0].ndim))

        argtys_extra = argtys
        sig_extra = ""
        result = None
        if 'out' in kwtys:
            argtys_extra += (kwtys['out'],)
            sig_extra += ", out=None"
            result = kwtys['out']

        if 'neighborhood' in kwtys:
            argtys_extra += (kwtys['neighborhood'],)
            sig_extra += ", neighborhood=None"

        # look in the type cache first
        if argtys_extra in self._type_cache:
            (_sig, _, _, _) = self._type_cache[argtys_extra]
            return _sig

        (real_ret, typemap, calltypes) = self.get_return_type(argtys)
        sig = signature(real_ret, *argtys_extra)
        dummy_text = ("def __numba_dummy_stencil({}{}):\n    pass\n".format(
                        ",".join(self.kernel_ir.arg_names), sig_extra))
        exec_(dummy_text) in globals(), locals()
        dummy_func = eval("__numba_dummy_stencil")
        sig.pysig = utils.pysignature(dummy_func)
        self._targetctx.insert_func_defn([(self._lower_me, self, argtys_extra)])
        self._type_cache[argtys_extra] = (sig, result, typemap, calltypes)
        return sig

    def copy_ir_with_calltypes(self, ir, calltypes):
        """
        Create a copy of a given IR along with its calltype information.
        We need a copy of the calltypes because copy propagation applied
        to the copied IR will change the calltypes and make subsequent
        uses of the original IR invalid.
        """
        copy_calltypes = {}
        kernel_copy = ir.copy()
        kernel_copy.blocks = {}
        # For each block...
        for (block_label, block) in ir.blocks.items():
            new_block = copy.deepcopy(ir.blocks[block_label])
            new_block.body = []
            # For each statement in each block...
            for stmt in ir.blocks[block_label].body:
                # Copy the statement to the new copy of the kernel
                # and if the original statement is in the original
                # calltypes then add the type associated with this
                # statement to the calltypes copy.
                scopy = copy.deepcopy(stmt)
                new_block.body.append(scopy)
                if stmt in calltypes:
                    copy_calltypes[scopy] = calltypes[stmt]
            kernel_copy.blocks[block_label] = new_block
        return (kernel_copy, copy_calltypes)

    def _stencil_wrapper(self, result, sigret, return_type, typemap, calltypes, *args):
        # Overall approach:
        # 1) Construct a string containing a function definition for the stencil function
        #    that will execute the stencil kernel.  This function definition includes a
        #    unique stencil function name, the parameters to the stencil kernel, loop
        #    nests across the dimenions of the input array.  Those loop nests use the
        #    computed stencil kernel size so as not to try to compute elements where
        #    elements outside the bounds of the input array would be needed.
        # 2) The but of the loop nest in this new function is a special sentinel
        #    assignment.
        # 3) Get the IR of this new function.
        # 4) Split the block containing the sentinel assignment and remove the sentinel
        #    assignment.  Insert the stencil kernel IR into the stencil function IR
        #    after label and variable renaming of the stencil kernel IR to prevent
        #    conflicts with the stencil function IR.
        # 5) Compile the combined stencil function IR + stencil kernel IR into existence.

        # Copy the kernel so that our changes for this callsite
        # won't effect other callsites.
        (kernel_copy, copy_calltypes) = self.copy_ir_with_calltypes(
                                            self.kernel_ir, calltypes)
        # The stencil kernel body becomes the body of a loop, for which args aren't needed.
        ir_utils.remove_args(kernel_copy.blocks)
        first_arg = kernel_copy.arg_names[0]

        in_cps, out_cps = ir_utils.copy_propagate(kernel_copy.blocks, typemap)
        name_var_table = ir_utils.get_name_var_table(kernel_copy.blocks)
        ir_utils.apply_copy_propagate(
            kernel_copy.blocks,
            in_cps,
            name_var_table,
            typemap,
            copy_calltypes)

        if "out" in name_var_table:
            raise ValueError("Cannot use the reserved word 'out' in stencil kernels.")

        sentinel_name = ir_utils.get_unused_var_name("__sentinel__", name_var_table)
        if config.DEBUG_ARRAY_OPT == 1:
            print("name_var_table", name_var_table, sentinel_name)

        the_array = args[0]

        if config.DEBUG_ARRAY_OPT == 1:
            print("_stencil_wrapper", return_type, return_type.dtype,
                                      type(return_type.dtype), args)
            ir_utils.dump_blocks(kernel_copy.blocks)

        # We generate a Numba function to execute this stencil and here
        # create the unique name of this function.
        stencil_func_name = "__numba_stencil_%s_%s" % (
                                        hex(id(the_array)).replace("-", "_"),
                                        self.id)

        # We will put a loop nest in the generated function for each
        # dimension in the input array.  Here we create the name for
        # the index variable for each dimension.  index0, index1, ...
        index_vars = []
        for i in range(the_array.ndim):
            index_var_name = ir_utils.get_unused_var_name("index" + str(i),
                                                          name_var_table)
            index_vars += [index_var_name]

        # Create extra signature for out and neighborhood.
        out_name = ir_utils.get_unused_var_name("out", name_var_table)
        neighborhood_name = ir_utils.get_unused_var_name("neighborhood",
                                                         name_var_table)
        sig_extra = ""
        if result is not None:
            sig_extra += ", {}=None".format(out_name)
        if "neighborhood" in dict(self.kws):
            sig_extra += ", {}=None".format(neighborhood_name)

        # Get a list of the standard indexed array names.
        standard_indexed = self.options.get("standard_indexing", [])

        if first_arg in standard_indexed:
            raise ValueError("The first argument to a stencil kernel must "
                             "use relative indexing, not standard indexing.")

        if len(set(standard_indexed) - set(kernel_copy.arg_names)) != 0:
            raise ValueError("Standard indexing requested for an array name "
                             "not present in the stencil kernel definition.")

        # Add index variables to getitems in the IR to transition the accesses
        # in the kernel from relative to regular Python indexing.  Returns the
        # computed size of the stencil kernel and a list of the relatively indexed
        # arrays.
        kernel_size, relatively_indexed = self.add_indices_to_kernel(
                kernel_copy, index_vars, the_array.ndim,
                self.neighborhood, standard_indexed)
        if self.neighborhood is None:
            self.neighborhood = kernel_size

        if config.DEBUG_ARRAY_OPT == 1:
            print("After add_indices_to_kernel")
            ir_utils.dump_blocks(kernel_copy.blocks)

        # The return in the stencil kernel becomes a setitem for that
        # particular point in the iteration space.
        ret_blocks = self.replace_return_with_setitem(kernel_copy.blocks,
                                                      index_vars, out_name)

        if config.DEBUG_ARRAY_OPT == 1:
            print("After replace_return_with_setitem", ret_blocks)
            ir_utils.dump_blocks(kernel_copy.blocks)

        # Start to form the new function to execute the stencil kernel.
        func_text = "def {}({}{}):\n".format(stencil_func_name,
                        ",".join(kernel_copy.arg_names), sig_extra)

        # Get loop ranges for each dimension, which could be either int
        # or variable. In the latter case we'll use the extra neighborhood
        # argument to the function.
        ranges = []
        for i in range(the_array.ndim):
            if isinstance(kernel_size[i][0], int):
                lo = kernel_size[i][0]
                hi = kernel_size[i][1]
            else:
                lo = "{}[{}][0]".format(neighborhood_name, i)
                hi = "{}[{}][1]".format(neighborhood_name, i)
            ranges.append((lo, hi))

        # If there are more than one relatively indexed arrays, add a call to
        # a function that will raise an error if any of the relatively indexed
        # arrays are of different size than the first input array.
        if len(relatively_indexed) > 1:
            func_text += "    raise_if_incompatible_array_sizes(" + first_arg
            for other_array in relatively_indexed:
                if other_array != first_arg:
                    func_text += "," + other_array
            func_text += ")\n"

        # Get the shape of the first input array.
        shape_name = ir_utils.get_unused_var_name("full_shape", name_var_table)
        func_text += "    {} = {}.shape\n".format(shape_name, first_arg)


        # If we have to allocate the output array (the out argument was not used)
        # then us numpy.full if the user specified a cval stencil decorator option
        # or np.zeros if they didn't to allocate the array.
        if result is None:
            if "cval" in self.options:
                cval = self.options["cval"]
                if return_type.dtype != typing.typeof.typeof(cval):
                    raise ValueError(
                        "cval type does not match stencil return type.")

                out_init ="{} = np.full({}, {}, dtype=np.{})\n".format(
                            out_name, shape_name, cval, return_type.dtype)

            else:
                out_init ="{} = np.zeros({}, dtype=np.{})\n".format(
                            out_name, shape_name, return_type.dtype)
            func_text += "    " + out_init

        offset = 1
        # Add the loop nests to the new function.
        for i in range(the_array.ndim):
            for j in range(offset):
                func_text += "    "
            # ranges[i][0] is the minimum index used in the i'th dimension
            # but minimum's greater than 0 don't preclude any entry in the array.
            # So, take the minimum of 0 and the minimum index found in the kernel
            # and this will be a negative number (potentially -0).  Then, we do
            # unary - on that to get the positive offset in this dimension whose
            # use is precluded.
            # ranges[i][1] is the maximum of 0 and the observed maximum index
            # in this dimension because negative maximums would not cause us to
            # preclude any entry in the array from being used.
            func_text += ("for {} in range(-min(0,{}),"
                          "{}[{}]-max(0,{})):\n").format(
                            index_vars[i],
                            ranges[i][0],
                            shape_name,
                            i,
                            ranges[i][1])
            offset += 1

        for j in range(offset):
            func_text += "    "
        # Put a sentinel in the code so we can locate it in the IR.  We will
        # remove this sentinel assignment and replace it with the IR for the
        # stencil kernel body.
        func_text += "{} = 0\n".format(sentinel_name)
        func_text += "    return {}\n".format(out_name)

        if config.DEBUG_ARRAY_OPT == 1:
            print("new stencil func text")
            print(func_text)

        # Force the new stencil function into existence.
        exec_(func_text) in globals(), locals()
        stencil_func = eval(stencil_func_name)
        if sigret is not None:
            pysig = utils.pysignature(stencil_func)
            sigret.pysig = pysig
        # Get the IR for the newly created stencil function.
        stencil_ir = compiler.run_frontend(stencil_func)
        ir_utils.remove_dels(stencil_ir.blocks)

        # rename all variables in stencil_ir afresh
        var_table = ir_utils.get_name_var_table(stencil_ir.blocks)
        new_var_dict = {}
        reserved_names = ([sentinel_name, out_name, neighborhood_name,
                           shape_name] + kernel_copy.arg_names + index_vars)
        for name, var in var_table.items():
            if not name in reserved_names:
                new_var_dict[name] = ir_utils.mk_unique_var(name)
        ir_utils.replace_var_names(stencil_ir.blocks, new_var_dict)

        stencil_stub_last_label = max(stencil_ir.blocks.keys()) + 1

        # Shift lables in the kernel copy so they are guaranteed unique
        # and don't conflict with any labels in the stencil_ir.
        kernel_copy.blocks = ir_utils.add_offset_to_labels(
                                kernel_copy.blocks, stencil_stub_last_label)
        new_label = max(kernel_copy.blocks.keys()) + 1
        # Adjust ret_blocks to account for addition of the offset.
        ret_blocks = [x + stencil_stub_last_label for x in ret_blocks]

        if config.DEBUG_ARRAY_OPT == 1:
            print("ret_blocks w/ offsets", ret_blocks, stencil_stub_last_label)
            print("before replace sentinel stencil_ir")
            ir_utils.dump_blocks(stencil_ir.blocks)
            print("before replace sentinel kernel_copy")
            ir_utils.dump_blocks(kernel_copy.blocks)

        # Search all the block in the stencil outline for the sentinel.
        for label, block in stencil_ir.blocks.items():
            for i, inst in enumerate(block.body):
                if (isinstance( inst, ir.Assign) and
                    inst.target.name == sentinel_name):
                    # We found the sentinel assignment.
                    loc = inst.loc
                    scope = block.scope
                    # split block across __sentinel__
                    # A new block is allocated for the statements prior to the
                    # sentinel but the new block maintains the current block
                    # label.
                    prev_block = ir.Block(scope, loc)
                    prev_block.body = block.body[:i]
                    # The current block is used for statements after sentinel.
                    block.body = block.body[i + 1:]
                    # But the current block gets a new label.
                    body_first_label = min(kernel_copy.blocks.keys())

                    # The previous block jumps to the minimum labelled block of
                    # the parfor body.
                    prev_block.append(ir.Jump(body_first_label, loc))
                    # Add all the parfor loop body blocks to the gufunc
                    # function's IR.
                    for (l, b) in kernel_copy.blocks.items():
                        stencil_ir.blocks[l] = b

                    stencil_ir.blocks[new_label] = block
                    stencil_ir.blocks[label] = prev_block
                    # Add a jump from all the blocks that previously contained
                    # a return in the stencil kernel to the block
                    # containing statements after the sentinel.
                    for ret_block in ret_blocks:
                        stencil_ir.blocks[ret_block].append(
                            ir.Jump(new_label, loc))
                    break
            else:
                continue
            break

        stencil_ir.blocks = ir_utils.rename_labels(stencil_ir.blocks)
        ir_utils.remove_dels(stencil_ir.blocks)

        assert(isinstance(the_array, types.Type))
        array_types = args

        new_stencil_param_types = list(array_types)

        if config.DEBUG_ARRAY_OPT == 1:
            print("new_stencil_param_types", new_stencil_param_types)
            ir_utils.dump_blocks(stencil_ir.blocks)

        # Compile the combined stencil function with the replaced loop
        # body in it.
        new_func = compiler.compile_ir(
            self._typingctx,
            self._targetctx,
            stencil_ir,
            new_stencil_param_types,
            None,
            compiler.DEFAULT_FLAGS,
            {})
        return new_func

    def __call__(self, *args, **kwargs):
        if (self.neighborhood is not None and
            len(self.neighborhood) != args[0].ndim):
            raise ValueError("{} dimensional neighborhood specified for {} "
                             "dimensional input array".format(
                                len(self.neighborhood), args[0].ndim))

        if 'out' in kwargs:
            result = kwargs['out']
            rdtype = result.dtype
            rttype = numpy_support.from_dtype(rdtype)
            result_type = types.npytypes.Array(rttype, result.ndim,
                                               numpy_support.map_layout(result))
            array_types = tuple([typing.typeof.typeof(x) for x in args])
            array_types_full = tuple([typing.typeof.typeof(x) for x in args] +
                                     [result_type])
        else:
            result = None
            array_types = tuple([typing.typeof.typeof(x) for x in args])
            array_types_full = array_types

        if config.DEBUG_ARRAY_OPT == 1:
            print("__call__", array_types, args, kwargs)

        (real_ret, typemap, calltypes) = self.get_return_type(array_types)
        new_func = self._stencil_wrapper(result, None, real_ret, typemap,
                                         calltypes, *array_types_full)

        if result is None:
            return new_func.entry_point(*args)
        else:
            return new_func.entry_point(*(args+(result,)))

def stencil(func_or_mode='constant', **options):
    # called on function without specifying mode style
    if not isinstance(func_or_mode, str):
        mode = 'constant'  # default style
        func = func_or_mode
    else:
        mode = func_or_mode
        func = None

    for option in options:
        if option not in ["cval", "standard_indexing", "neighborhood"]:
            raise ValueError("Unknown stencil option " + option)

    wrapper = _stencil(mode, options)
    if func is not None:
        return wrapper(func)
    return wrapper

def _stencil(mode, options):
    if mode != 'constant':
        raise ValueError("Unsupported mode style " + mode)

    def decorated(func):
        kernel_ir = compiler.run_frontend(func)
        return StencilFunc(kernel_ir, mode, options)

    return decorated

@lower_builtin(stencil)
def stencil_dummy_lower(context, builder, sig, args):
    "lowering for dummy stencil calls"
    return lir.Constant(lir.IntType(types.intp.bitwidth), 0)