How to use the numba.ir.Assign 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 IntelPython / sdc / sdc / distributed.py View on Github external
# str_out = df2.to_csv(None, header=header)
            # sdc.io.np_io._file_write_parallel(fname, str_out)

            df_typ = self.state.typemap[df.name]
            rhs = assign.value
            fname = args[0]

            # update df index and get to_csv from new df
            nodes = self._fix_parallel_df_index(df)
            new_df = nodes[-1].target
            new_df_typ = self.state.typemap[new_df.name]
            new_to_csv = ir.Expr.getattr(new_df, 'to_csv', new_df.loc)
            new_func = ir.Var(new_df.scope, mk_unique_var('func'), new_df.loc)
            self.state.typemap[new_func.name] = self.state.typingctx.resolve_getattr(
                new_df_typ, 'to_csv')
            nodes.append(ir.Assign(new_to_csv, new_func, new_df.loc))
            rhs.func = new_func

            # # header = header and is_root
            kws = dict(rhs.kws)
            true_var = ir.Var(assign.target.scope, mk_unique_var('true'), rhs.loc)
            self.state.typemap[true_var.name] = types.bool_
            nodes.append(
                ir.Assign(ir.Const(True, new_df.loc), true_var, new_df.loc))
            header_var = self._get_arg(
                'to_csv', rhs.args, kws, 5, 'header', true_var)
            nodes += self._gen_is_root_and_cond(header_var)
            header_var = nodes[-1].target
            if len(rhs.args) > 5:
                rhs.args[5] = header_var
            else:
                kws['header'] = header_var
github IntelPython / sdc / sdc / hiframes / hiframes_typed.py View on Github external
n_arg = rhs.args[0]
            func = series_replace_funcs[func_name]
            return self._replace_func(func, [data, n_arg, name],
                                      extra_globals={'gt_f': series_kernels.gt_f,
                                                     'lt_f': series_kernels.lt_f},
                                      pre_nodes=nodes)

        if func_name == 'head':
            nodes = []
            n_arg = self._get_arg('Series.head', rhs.args, dict(rhs.kws), 0,
                                  'n', default=False)  # TODO: proper default handling
            if n_arg is False:
                n_arg = ir.Var(lhs.scope, mk_unique_var('head_n'), lhs.loc)
                # default is 5
                self.state.typemap[n_arg.name] = types.IntegerLiteral(5)
                nodes.append(ir.Assign(
                    ir.Const(5, lhs.loc), n_arg, lhs.loc))

            data = self._get_series_data(series_var, nodes)
            func = series_replace_funcs[func_name]

            if self.state.typemap[series_var.name].index != types.none:
                index = self._get_series_index(series_var, nodes)
                func = series_replace_funcs['head_index']
            else:
                index = self._get_index_values(data, nodes)

            name = self._get_series_name(series_var, nodes)

            return self._replace_func(
                func, (data, index, n_arg, name), pre_nodes=nodes)
github numba / numba / numba / stencilparfor.py View on Github external
stencil_blocks[parfor_body_exit_label].body.append(
            ir.Return(ret_const_var, dummy_loc),
        )
        stencil_blocks = ir_utils.simplify_CFG(stencil_blocks)
        stencil_blocks[max(stencil_blocks.keys())].body.pop()

        if config.DEBUG_ARRAY_OPT >= 1:
            print("stencil_blocks after adding SetItem")
            ir_utils.dump_blocks(stencil_blocks)

        pattern = ('stencil', [start_lengths, end_lengths])
        parfor = numba.parfor.Parfor(loopnests, init_block, stencil_blocks,
                                     loc, parfor_ind_var, equiv_set, pattern, self.flags)
        gen_nodes.append(parfor)
        gen_nodes.append(ir.Assign(out_arr, target, loc))
        return gen_nodes
github numba / numba / numba / parfor.py View on Github external
if call_name is 'dot':
                # dot has two inputs
                tmp_var1 = tmp_var
                in2 = args[1]
                tmp_var2 = ir.Var(scope, mk_unique_var("$val"), loc)
                self.typemap[tmp_var2.name] = in_typ
                getitem_call2 = ir.Expr.getitem(in2, index_var, loc)
                self.calltypes[getitem_call2] = signature(
                    in_typ, arr_typ, index_var_type)
                acc_block.body.append(ir.Assign(getitem_call2, tmp_var2, loc))
                mult_call = ir.Expr.binop('*', tmp_var1, tmp_var2, loc)
                mult_func_typ = find_op_typ('*', [in_typ, in_typ])
                self.calltypes[mult_call] = mult_func_typ
                tmp_var = ir.Var(scope, mk_unique_var("$val"), loc)
                acc_block.body.append(ir.Assign(mult_call, tmp_var, loc))

            acc_call = ir.Expr.inplace_binop(
                acc_op, im_op, acc_var, tmp_var, loc)
            # for some reason, type template of += returns None,
            # so type template of + should be used
            self.calltypes[acc_call] = im_op_func_typ
            # FIXME: we had to break assignment: acc += ... acc ...
            # into two assignment: acc_tmp = ... acc ...; x = acc_tmp
            # in order to avoid an issue in copy propagation.
            acc_tmp_var = ir.Var(scope, mk_unique_var("$acc"), loc)
            self.typemap[acc_tmp_var.name] = el_typ
            acc_block.body.append(ir.Assign(acc_call, acc_tmp_var, loc))
            acc_block.body.append(ir.Assign(acc_tmp_var, acc_var, loc))
            loop_body = {next_label(): acc_block}

            # parfor
github numba / numba / numba / npyufunc / parfor.py View on Github external
if isinstance(redvar_typ, types.npytypes.Array):
                    # Create an array of identity values for the reduction.
                    # First, create a variable for np.full.
                    full_func = ir.Var(scope, mk_unique_var("full_func"), loc)
                    full_fnty = get_np_ufunc_typ(np.full)
                    full_sig = typingctx.resolve_function_type(full_fnty,
                               (types.UniTuple(types.intp, redvar_typ.ndim),
                                reddtype,
                                types.DType(reddtype)), {})
                    typemap[full_func.name] = full_fnty
                    full_assign = ir.Assign(ir.Global("full", np.full, loc=loc), full_func, loc)
                    lowerer.lower_inst(full_assign)

                    # Then create a var with the identify value.
                    init_val_var = ir.Var(scope, mk_unique_var("init_val"), loc)
                    init_val_assign = ir.Assign(ir.Const(init_val, loc), init_val_var, loc)
                    typemap[init_val_var.name] = reddtype
                    lowerer.lower_inst(init_val_assign)

                    # Then, call np.full with the shape of the reduction array and the identity value.
                    full_call = ir.Expr.call(full_func, [redshape_var, init_val_var], {}, loc=loc)
                    lowerer.fndesc.calltypes[full_call] = full_sig
                    redtoset = ir.Var(scope, mk_unique_var("redtoset"), loc)
                    redtoset_assign = ir.Assign(full_call, redtoset, loc)
                    typemap[redtoset.name] = redvar_typ
                    lowerer.lower_inst(redtoset_assign)
                else:
                    redtoset = ir.Var(scope, mk_unique_var("redtoset"), loc)
                    redtoset_assign = ir.Assign(ir.Const(init_val, loc), redtoset, loc)
                    typemap[redtoset.name] = reddtype
                    lowerer.lower_inst(redtoset_assign)
            else:
github numba / numba / numba / stencil.py View on Github external
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
github numba / numba / numba / parfor.py View on Github external
arg_vars.append(arg_out_var)
        if op in npydecl.supported_array_operators:
            el_typ1 = typemap[arg_vars[0].name]
            if len(arg_vars) == 2:
                el_typ2 = typemap[arg_vars[1].name]
                func_typ = find_op_typ(op, [el_typ1, el_typ2])
                ir_expr = ir.Expr.binop(op, arg_vars[0], arg_vars[1], loc)
                if op == '/':
                    func_typ, ir_expr = _gen_np_divide(
                        arg_vars[0], arg_vars[1], out_ir, typemap)
            else:
                func_typ = find_op_typ(op, [el_typ1])
                ir_expr = ir.Expr.unary(op, arg_vars[0], loc)
            calltypes[ir_expr] = func_typ
            el_typ = func_typ.return_type
            out_ir.append(ir.Assign(ir_expr, expr_out_var, loc))
        for T in array_analysis.MAP_TYPES:
            if isinstance(op, T):
                # elif isinstance(op, (np.ufunc, DUFunc)):
                # function calls are stored in variables which are not removed
                # op is typing_key to the variables type
                func_var_name = _find_func_var(typemap, op, avail_vars)
                func_var = ir.Var(scope, mk_unique_var(func_var_name), loc)
                typemap[func_var.name] = typemap[func_var_name]
                func_var_def = func_ir.get_definition(func_var_name)
                if isinstance(func_var_def, ir.Expr) and func_var_def.op == 'getattr' and func_var_def.attr == 'sqrt':
                     g_math_var = ir.Var(scope, mk_unique_var("$math_g_var"), loc)
                     typemap[g_math_var.name] = types.misc.Module(math)
                     g_math = ir.Global('math', math, loc)
                     g_math_assign = ir.Assign(g_math, g_math_var, loc)
                     func_var_def = ir.Expr.getattr(g_math_var, 'sqrt', loc)
                     out_ir.append(g_math_assign)
github numba / numba / numba / stencilparfor.py View on Github external
init_block.body.extend(stmts)
        else: # out is present
            if "cval" in stencil_func.options: # do out[:] = cval
                cval = stencil_func.options["cval"]
                # TODO: Loosen this restriction to adhere to casting rules.
                cval_ty = typing.typeof.typeof(cval)
                if not self.typingctx.can_convert(cval_ty, return_type.dtype):
                    msg = "cval type does not match stencil return type."
                    raise ValueError(msg)

                # get slice ref
                slice_var = ir.Var(scope, mk_unique_var("$py_g_var"), loc)
                slice_fn_ty = self.typingctx.resolve_value_type(slice)
                self.typemap[slice_var.name] = slice_fn_ty
                slice_g = ir.Global('slice', slice, loc)
                slice_assigned = ir.Assign(slice_g, slice_var, loc)
                init_block.body.append(slice_assigned)

                sig = self.typingctx.resolve_function_type(slice_fn_ty,
                                                           (types.none,) * 2,
                                                           {})

                callexpr = ir.Expr.call(func=slice_var, args=(), kws=(),
                                        loc=loc)

                self.calltypes[callexpr] = sig
                slice_inst_var = ir.Var(scope, mk_unique_var("$slice_inst"),
                                        loc)
                self.typemap[slice_inst_var.name] = types.slice2_type
                slice_assign = ir.Assign(callexpr, slice_inst_var, loc)
                init_block.body.append(slice_assign)
github IntelPython / sdc / sdc / distributed.py View on Github external
rhs.args[3] = true_var
            out = [ir.Assign(ir.Const(True, loc), true_var, loc), assign]

        if fdef == ('rolling_variable', 'sdc.hiframes.rolling') and (
                self._is_1D_arr(rhs.args[0].name)
                or self._is_1D_Var_arr(rhs.args[0].name)):
            in_arr = rhs.args[0].name
            if self._is_1D_arr(in_arr):
                self._array_starts[lhs] = self._array_starts[in_arr]
                self._array_counts[lhs] = self._array_counts[in_arr]
                self._array_sizes[lhs] = self._array_sizes[in_arr]
            # set parallel flag to true
            true_var = ir.Var(scope, mk_unique_var("true_var"), loc)
            self.state.typemap[true_var.name] = types.boolean
            rhs.args[4] = true_var
            out = [ir.Assign(ir.Const(True, loc), true_var, loc), assign]

        if (func_mod == 'sdc.hiframes.rolling'
            and func_name in ('shift', 'pct_change')
            and (self._is_1D_arr(rhs.args[0].name)
                 or self._is_1D_Var_arr(rhs.args[0].name))):
            in_arr = rhs.args[0].name
            if self._is_1D_arr(in_arr):
                self._array_starts[lhs] = self._array_starts[in_arr]
                self._array_counts[lhs] = self._array_counts[in_arr]
                self._array_sizes[lhs] = self._array_sizes[in_arr]
            # set parallel flag to true
            true_var = ir.Var(scope, mk_unique_var("true_var"), loc)
            self.state.typemap[true_var.name] = types.boolean
            rhs.args[2] = true_var
            out = [ir.Assign(ir.Const(True, loc), true_var, loc), assign]
github IntelPython / sdc / sdc / hiframes / hiframes_typed.py View on Github external
def _run_unary(self, assign, rhs):
        arg = rhs.value
        typ = self.state.typemap[arg.name]

        if isinstance(typ, SeriesType):
            nodes = []
            arg = self._get_series_data(arg, nodes)
            rhs.value = arg
            self._convert_series_calltype(rhs)
            out_data = ir.Var(
                arg.scope, mk_unique_var(assign.target.name + '_data'), rhs.loc)
            self.state.typemap[out_data.name] = self.state.calltypes[rhs].return_type
            nodes.append(ir.Assign(rhs, out_data, rhs.loc))
            return self._replace_func(
                lambda data: sdc.hiframes.api.init_series(data),
                [out_data],
                pre_nodes=nodes
            )

        return [assign]