How to use the numba.ir_utils.guard 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 / stencilparfor.py View on Github external
cases.
    """
    require(isinstance(index_var, ir.Var))
    # case where the index is a const itself in outer function
    var_const =  guard(_get_const_two_irs, stencil_ir, func_ir, index_var)
    if var_const is not None:
        return var_const
    # get index definition
    index_def = ir_utils.get_definition(stencil_ir, index_var)
    # match inner_var = unary(index_var)
    var_const = guard(
        _get_const_unary_expr, stencil_ir, func_ir, index_def)
    if var_const is not None:
        return var_const
    # match inner_var = arg1 + arg2
    var_const = guard(
        _get_const_binary_expr, stencil_ir, func_ir, index_def)
    if var_const is not None:
        return var_const
    raise GuardException
github numba / numba / numba / parfor.py View on Github external
topo_order = find_topo_order(blocks)
        for label in topo_order:
            block = blocks[label]
            new_body = []
            equiv_set = self.array_analysis.get_equiv_set(label)
            for instr in block.body:
                parfor = None
                if isinstance(instr, ir.Assign):
                    loc = instr.loc
                    lhs = instr.target
                    expr = instr.value
                    callname = guard(find_callname, self.func_ir, expr)
                    if (callname == ('reduce', 'builtins')
                        or callname == ('reduce', '_functools')):
                        # reduce function with generic function
                        parfor = guard(self._reduce_to_parfor, equiv_set, lhs,
                                       expr.args, loc)
                    if parfor:
                        instr = parfor
                new_body.append(instr)
            block.body = new_body
        return
github IntelPython / sdc / sdc / hiframes / hiframes_typed.py View on Github external
from numba.stencil import StencilFunc
            # could be make_function from list comprehension which is ok
            func_def = guard(get_definition, self.state.func_ir, rhs.func)
            if isinstance(func_def, ir.Expr) and func_def.op == 'make_function':
                return [assign]
            if isinstance(func_def, ir.Global) and isinstance(func_def.value, StencilFunc):
                return [assign]
            warnings.warn(
                "function call couldn't be found for initial analysis")
            return [assign]
        else:
            func_name, func_mod = fdef

        string_methods_types = (SplitViewStringMethodsType, StringMethodsType)
        if isinstance(func_mod, ir.Var) and isinstance(self.state.typemap[func_mod.name], string_methods_types):
            f_def = guard(get_definition, self.state.func_ir, rhs.func)
            str_def = guard(get_definition, self.state.func_ir, f_def.value)
            if str_def is None:  # TODO: check for errors
                raise ValueError("invalid series.str")

            series_var = str_def.value

            # functions which are used from Numba directly by calling from StringMethodsType
            # other functions (for example, 'capitalize' is not presented in Numba) goes to be replaced here
            if func_name not in sdc.hiframes.pd_series_ext.str2str_methods_excluded:
                return self._run_series_str_method(assign, assign.target, series_var, func_name, rhs)

        # replace _get_type_max_value(arr.dtype) since parfors
        # arr.dtype transformation produces invalid code for dt64
        # TODO: min
        if fdef == ('_get_type_max_value', 'sdc.hiframes.hiframes_typed'):
            if self.state.typemap[rhs.args[0].name] == types.DType(types.NPDatetime('ns')):
github IntelPython / sdc / sdc / hiframes / hiframes_typed.py View on Github external
def _run_series_str_replace(self, assign, lhs, arr, rhs, nodes):
        regex = True
        # TODO: refactor arg parsing
        kws = dict(rhs.kws)
        if 'regex' in kws:
            regex = guard(find_const, self.state.func_ir, kws['regex'])
            if regex is None:
                raise ValueError(
                    "str.replace regex argument should be constant")

        impl = (series_kernels._str_replace_regex_impl if regex
                else series_kernels._str_replace_noregex_impl)

        return self._replace_func(
            impl,
            [arr, rhs.args[0], rhs.args[1]], pre_nodes=nodes,
            extra_globals={'unicode_to_std_str': unicode_to_std_str,
                           'std_str_to_unicode': std_str_to_unicode,
                           'pre_alloc_string_array': pre_alloc_string_array,
                           'get_utf8_size': get_utf8_size,
                           're': re}
        )
github IntelPython / sdc / sdc / hiframes / dataframe_pass.py View on Github external
f = loc_vars['f']

        f_ir = compile_to_numba_ir(
            f, {'numba': numba, 'np': np, 'Row': Row, 'sdc': sdc})
        # fix definitions to enable finding sentinel
        f_ir._definitions = build_definitions(f_ir.blocks)
        topo_order = find_topo_order(f_ir.blocks)

        # find sentinel function and replace with user func
        for l in topo_order:
            block = f_ir.blocks[l]
            for i, stmt in enumerate(block.body):
                if (isinstance(stmt, ir.Assign)
                        and isinstance(stmt.value, ir.Expr)
                        and stmt.value.op == 'call'):
                    fdef = guard(get_definition, f_ir, stmt.value.func)
                    if isinstance(fdef, ir.Global) and fdef.name == 'map_func':
                        update_globals(func, _globals)
                        inline_closure_call(f_ir, _globals, block, i, func)
                        # fix the global value to avoid typing errors
                        fdef.value = 1
                        fdef.name = 'A'
                        break

        arg_typs = tuple(df_typ.data[df_typ.columns.index(c)] for c in used_cols)
        f_typemap, f_return_type, f_calltypes = numba.typed_passes.type_inference_stage(
            self.state.typingctx, f_ir, arg_typs, None)
        self.state.typemap.update(f_typemap)
        self.state.calltypes.update(f_calltypes)

        nodes = []
        col_vars = [self._get_dataframe_data(df_var, c, nodes) for c in used_cols]
github numba / numba / numba / parfor.py View on Github external
loc = instr.loc
                    target = instr.target
                    index = instr.index if isinstance(instr, ir.SetItem) else instr.index_var
                    value = instr.value
                    target_typ = self.typemap[target.name]
                    index_typ = self.typemap[index.name]
                    value_typ = self.typemap[value.name]
                    if isinstance(target_typ, types.npytypes.Array):
                        if (isinstance(index_typ, types.npytypes.Array) and
                            isinstance(index_typ.dtype, types.Boolean) and
                            target_typ.ndim == index_typ.ndim):
                            if isinstance(value_typ, types.Number):
                                instr = self._setitem_to_parfor(equiv_set,
                                        loc, target, index, value)
                            elif isinstance(value_typ, types.npytypes.Array):
                                val_def = guard(get_definition, self.func_ir,
                                                value.name)
                                if (isinstance(val_def, ir.Expr) and
                                    val_def.op == 'getitem' and
                                    val_def.index.name == index.name):
                                    instr = self._setitem_to_parfor(equiv_set,
                                            loc, target, index, val_def.value)
                        else:
                            shape = equiv_set.get_shape(instr)
                            if shape != None:
                                instr = self._setitem_to_parfor(equiv_set,
                                        loc, target, index, value, shape=shape)
                new_body.append(instr)
            block.body = new_body
github IntelPython / sdc / sdc / hiframes / dataframe_pass.py View on Github external
def _get_dataframe_data(self, df_var, col_name, nodes):
        # optimization: return data var directly if not ambiguous
        # (no multiple init_dataframe calls for the same df_var with control
        # flow)
        # e.g. A = init_dataframe(A, None, 'A')
        # XXX assuming init_dataframe is the only call to create a dataframe
        # and dataframe._data is never overwritten
        df_typ = self.state.typemap[df_var.name]
        ind = df_typ.columns.index(col_name)
        var_def = guard(get_definition, self.state.func_ir, df_var)
        call_def = guard(find_callname, self.state.func_ir, var_def)
        if call_def == ('init_dataframe', 'sdc.hiframes.pd_dataframe_ext'):
            return var_def.args[ind]

        loc = df_var.loc
        ind_var = ir.Var(df_var.scope, mk_unique_var('col_ind'), loc)
        self.state.typemap[ind_var.name] = types.IntegerLiteral(ind)
        nodes.append(ir.Assign(ir.Const(ind, loc), ind_var, loc))
        # XXX use get_series_data() for getting data instead of S._data
        # to enable alias analysis
        f_block = compile_to_numba_ir(
            lambda df, c_ind: sdc.hiframes.pd_dataframe_ext.get_dataframe_data(
                df, c_ind),
            {'sdc': sdc},
            self.state.typingctx,
            (df_typ, self.state.typemap[ind_var.name]),
github IntelPython / sdc / sdc / hiframes / dataframe_pass.py View on Github external
def _get_dataframe_index(self, df_var, nodes):
        df_typ = self.state.typemap[df_var.name]
        n_cols = len(df_typ.columns)
        var_def = guard(get_definition, self.state.func_ir, df_var)
        call_def = guard(find_callname, self.state.func_ir, var_def)
        if call_def == ('init_dataframe', 'sdc.hiframes.pd_dataframe_ext'):
            return var_def.args[n_cols]

        # XXX use get_series_data() for getting data instead of S._data
        # to enable alias analysis
        f_block = compile_to_numba_ir(
            lambda df: sdc.hiframes.pd_dataframe_ext.get_dataframe_index(df),
            {'sdc': sdc},
            self.state.typingctx,
            (df_typ,),
            self.state.typemap,
            self.state.calltypes
        ).blocks.popitem()[1]
        replace_arg_nodes(f_block, [df_var])
        nodes += f_block.body[:-2]
        return nodes[-1].target
github IntelPython / sdc / sdc / hiframes / dataframe_pass.py View on Github external
def _get_dataframe_index(self, df_var, nodes):
        df_typ = self.state.typemap[df_var.name]
        n_cols = len(df_typ.columns)
        var_def = guard(get_definition, self.state.func_ir, df_var)
        call_def = guard(find_callname, self.state.func_ir, var_def)
        if call_def == ('init_dataframe', 'sdc.hiframes.pd_dataframe_ext'):
            return var_def.args[n_cols]

        # XXX use get_series_data() for getting data instead of S._data
        # to enable alias analysis
        f_block = compile_to_numba_ir(
            lambda df: sdc.hiframes.pd_dataframe_ext.get_dataframe_index(df),
            {'sdc': sdc},
            self.state.typingctx,
            (df_typ,),
            self.state.typemap,
            self.state.calltypes
        ).blocks.popitem()[1]
        replace_arg_nodes(f_block, [df_var])
        nodes += f_block.body[:-2]
        return nodes[-1].target
github IntelPython / sdc / sdc / hiframes / dataframe_pass.py View on Github external
def _run_setitem(self, inst):
        target_typ = self.state.typemap[inst.target.name]
        nodes = []
        index_var = (inst.index_var if isinstance(inst, ir.StaticSetItem)
                     else inst.index)
        index_typ = self.state.typemap[index_var.name]

        if self._is_df_iat_var(inst.target):
            df_var = guard(get_definition, self.state.func_ir, inst.target).value
            df_typ = self.state.typemap[df_var.name]
            val = inst.value
            # df.iat[n,1] = 3
            if isinstance(index_typ, types.Tuple) and len(index_typ) == 2:
                ind_def = guard(get_definition, self.state.func_ir, index_var)
                col_ind = guard(find_const, self.state.func_ir, ind_def.items[1])
                col_name = df_typ.columns[col_ind]
                in_arr = self._get_dataframe_data(df_var, col_name, nodes)
                row_ind = ind_def.items[0]

                def _impl(A, row_ind, val):
                    A[row_ind] = val
                return self._replace_func(_impl,
                                          [in_arr, row_ind, val], pre_nodes=nodes)
        return [inst]