How to use the numba.ir function in numba

To help you get started, we’ve selected a few numba examples, based on popular ways it is used in public projects.

Secure your code as it's written. Use Snyk Code to scan source code in minutes - no build needed - and fix issues immediately.

github numba / numba / numba / stencil.py View on Github external
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))
                            # Get the type of this particular part of the index tuple.
                            one_index_typ = stmt_index_var_typ[dim]
                            # If the array is indexed with a slice then we
                            # have to add the index value with a call to
                            # slice_addition.
                            if isinstance(one_index_typ, types.misc.SliceType):
                                sa_var = ir.Var(scope, ir_utils.mk_unique_var("slice_addition"), loc)
                                sa_func = numba.njit(slice_addition)
                                sa_func_typ = types.functions.Dispatcher(sa_func)
                                typemap[sa_var.name] = sa_func_typ
                                g_sa = ir.Global("slice_addition", sa_func, loc)
                                new_body.append(ir.Assign(g_sa, sa_var, loc))
                                slice_addition_call = ir.Expr.call(sa_var, [getitemvar, index_vars[dim]], (), loc)
                                calltypes[slice_addition_call] = sa_func_typ.get_call_type(self._typingctx, [one_index_typ, types.intp], {})
                                new_body.append(ir.Assign(slice_addition_call, tmpvar, loc))
                            else:
                                acc_call = ir.Expr.binop(operator.add, 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)
github numba / numba / numba / parfor.py View on Github external
in1 = args[0]
            in2 = args[1]
            el_typ = self.typemap[lhs.name].dtype
            assert self._get_ndims(
                in1.name) <= 2 and self._get_ndims(
                in2.name) == 1
            # loop range correlation is same as first dimention of 1st input
            corr = self.array_analysis.array_shape_classes[in1.name][0]
            size_var = self.array_analysis.array_size_vars[in1.name][0]
            scope = lhs.scope
            loc = expr.loc
            index_var = ir.Var(scope, mk_unique_var("parfor_index"), lhs.loc)
            self.typemap[index_var.name] = types.intp
            loopnests = [LoopNest(index_var, 0, size_var, 1, corr)]
            init_block = ir.Block(scope, loc)
            parfor = Parfor(
                loopnests,
                init_block,
                {},
                loc,
                self.array_analysis,
                index_var)
            if self._get_ndims(in1.name) == 2:
                # for 2D input, there is an inner loop
                # correlation of inner dimension
                inner_size_var = self.array_analysis.array_size_vars[in1.name][1]
                # loop structure: range block, header block, body

                range_label = next_label()
                header_label = next_label()
                body_label = next_label()
github numba / numba / numba / lowering.py View on Github external
def lower_assign(self, inst):
        """
        The returned object must have a new reference
        """
        value = inst.value
        if isinstance(value, ir.Const):
            return self.lower_const(value.value)
        elif isinstance(value, ir.Var):
            val = self.loadvar(value.name)
            self.incref(val)
            return val
        elif isinstance(value, ir.Expr):
            return self.lower_expr(value)
        elif isinstance(value, ir.Global):
            return self.lower_global(value.name, value.value)
        else:
            raise NotImplementedError(type(value), value)
github numba / numba / numba / parfor.py View on Github external
# make sure init block isn't removed
        init_block_label = min(blocks.keys())
        # update loop body blocks
        blocks.pop(init_block_label)
        parfor.loop_body = blocks

    # make sure dummy jump to loop body isn't altered
    first_body_label = min(parfor.loop_body.keys())
    assert isinstance(parfor.init_block.body[-1], ir.Jump)

    # remove dummy jump to loop body
    parfor.init_block.body.pop()

    # make sure dummy jump back to loop body isn't altered
    for block in parfor.loop_body.values():
        if (isinstance(block.body[-1], ir.Jump) and
            block.body[-1].target == first_body_label):
            # remove dummy jump back to loop
            block.body.pop()
    return
github numba / numba / numba / lowering.py View on Github external
def lower_inst(self, inst):
        if isinstance(inst, ir.Assign):
            value = self.lower_assign(inst)
            self.storevar(value, inst.target.name)

        elif isinstance(inst, ir.SetItem):
            target = self.loadvar(inst.target.name)
            index = self.loadvar(inst.index.name)
            value = self.loadvar(inst.value.name)
            ok = self.pyapi.object_setitem(target, index, value)
            negone = lc.Constant.int_signextend(ok.type, -1)
            pred = self.builder.icmp(lc.ICMP_EQ, ok, negone)
            with cgutils.if_unlikely(self.builder, pred):
                self.return_exception_raised()

        elif isinstance(inst, ir.Return):
            retval = self.loadvar(inst.value.name)
            self.incref(retval)
github numba / numba / numba / ir_utils.py View on Github external
for label in reversed(order):
        for inst in reversed(blocks[label].body):
            if isinstance(inst, ir.Assign):
                lhs = inst.target.name
                rhs = inst.value
                if isinstance(rhs, ir.Expr) and rhs.op == 'call':
                    call_table[rhs.func.name] = []
                if isinstance(rhs, ir.Expr) and rhs.op == 'getattr':
                    if lhs in call_table:
                        call_table[lhs].append(rhs.attr)
                        reverse_call_table[rhs.value.name] = lhs
                    if lhs in reverse_call_table:
                        call_var = reverse_call_table[lhs]
                        call_table[call_var].append(rhs.attr)
                        reverse_call_table[rhs.value.name] = call_var
                if isinstance(rhs, ir.Global):
                    if lhs in call_table:
                        call_table[lhs].append(rhs.value)
                    if lhs in reverse_call_table:
                        call_var = reverse_call_table[lhs]
                        call_table[call_var].append(rhs.value)
                if isinstance(rhs, ir.FreeVar):
                    if lhs in call_table:
                        call_table[lhs].append(rhs.value)
                    if lhs in reverse_call_table:
                        call_var = reverse_call_table[lhs]
                        call_table[call_var].append(rhs.value)
                if isinstance(rhs, ir.Var):
                    if lhs in call_table:
                        call_table[lhs].append(rhs.name)
                        reverse_call_table[rhs.name] = lhs
                    if lhs in reverse_call_table:
github IntelPython / sdc / hpat / hiframes / filter.py View on Github external
def filter_distributed_run(filter_node, array_dists, typemap, calltypes, typingctx, targetctx, dist_pass):
    # TODO: rebalance if output distributions are 1D instead of 1D_Var
    loc = filter_node.loc
    bool_arr = filter_node.bool_arr

    out = []
    for col_name, col_in_var in filter_node.df_in_vars.items():
        col_out_var = filter_node.df_out_vars[col_name]
        # using getitem like Numba for filtering arrays
        # TODO: generate parfor
        getitem_call = ir.Expr.getitem(col_in_var, bool_arr, loc)
        calltypes[getitem_call] = signature(
            typemap[col_out_var.name],  # output type
            typemap[col_in_var.name],  # input type
            typemap[bool_arr.name])  # index type
        out.append(ir.Assign(getitem_call, col_out_var, loc))

    return out
github numba / numba / numba / npyufunc / parfor.py View on Github external
def wrap_loop_body(loop_body):
    blocks = loop_body.copy()  # shallow copy is enough
    first_label = min(blocks.keys())
    last_label = max(blocks.keys())
    loc = blocks[last_label].loc
    blocks[last_label].body.append(ir.Jump(first_label, loc))
    return blocks
github numba / numba / numba / npyufunc / parfor.py View on Github external
ast_call = ast.Call(fn_ast_name, ast_args, [], None, None)
            return ast_call, env
    elif isinstance(expr, ir.Var):
        name = expr.name
        var = ast.Name(name, ast.Load(),
                        lineno=expr.loc.line,
                        col_offset=expr.loc.col if expr.loc.col else 0)
        typ = typemap[name]
        if isinstance(typ, types.Array):
            var = ast.Subscript(
                value = var,
                slice = ast.Index(value = _mk_tuple([ast.Name(v, ast.Load()) for v in subscripts])),
                ctx = ast.Load())
        return var, {}

    elif isinstance(expr, ir.Const):
        return ast.Num(expr.value), {}
    raise NotImplementedError(
        "Don't know how to translate array expression '%r'" % (expr,))
github IntelPython / sdc / sdc / distributed.py View on Github external
def _set_getsetitem_index(node, new_ind):
    if ((isinstance(node, ir.Expr) and node.op == 'static_getitem') or isinstance(node, ir.StaticSetItem)):
        node.index_var = new_ind
        node.index = None
        return

    assert ((isinstance(node, ir.Expr) and node.op == 'getitem') or isinstance(node, ir.SetItem))
    node.index = new_ind