def include_new_blocks(blocks, new_blocks, label, new_body, remove_non_return=True, work_list=None, func_ir=None): inner_blocks = add_offset_to_labels(new_blocks, ir_utils._max_label + 1) blocks.update(inner_blocks) ir_utils._max_label = max(blocks.keys()) scope = blocks[label].scope loc = blocks[label].loc inner_topo_order = find_topo_order(inner_blocks) inner_first_label = inner_topo_order[0] inner_last_label = inner_topo_order[-1] if remove_non_return: remove_return_from_block(inner_blocks[inner_last_label]) new_body.append(ir.Jump(inner_first_label, loc)) blocks[label].body = new_body label = ir_utils.next_label() blocks[label] = ir.Block(scope, loc) if remove_non_return: inner_blocks[inner_last_label].body.append(ir.Jump(label, loc)) # new_body.clear() if work_list is not None: topo_order = find_topo_order(inner_blocks) for _label in topo_order: block = inner_blocks[_label] block.scope = scope numba.inline_closurecall._add_definitions(func_ir, block) work_list.append((_label, block)) return label
def run(self): blocks = self.func_ir.blocks call_table, _ = ir_utils.get_call_table(blocks) topo_order = find_topo_order(blocks) for label in topo_order: new_body = [] for inst in blocks[label].body: if isinstance(inst, ir.Assign): out_nodes = self._run_assign(inst, call_table) if isinstance(out_nodes, list): new_body.extend(out_nodes) if isinstance(out_nodes, dict): label = include_new_blocks(blocks, out_nodes, label, new_body) new_body = [] if isinstance(out_nodes, tuple): gen_blocks, post_nodes = out_nodes label = include_new_blocks(blocks, gen_blocks, label, new_body) new_body = post_nodes else: new_body.append(inst) blocks[label].body = new_body self.func_ir._definitions = get_definitions(self.func_ir.blocks) return
def _run_pd_DatetimeIndex(self, assign, lhs, rhs): """transform pd.DatetimeIndex() call with string array argument """ kws = dict(rhs.kws) if 'data' in kws: data = kws['data'] if len(rhs.args) != 0: # pragma: no cover raise ValueError( "only data argument suppoted in pd.DatetimeIndex()") else: if len(rhs.args) != 1: # pragma: no cover raise ValueError( "data argument in pd.DatetimeIndex() expected") data = rhs.args[0] def f(str_arr): numba.parfor.init_prange() n = len(str_arr) S = numba.unsafe.ndarray.empty_inferred((n, )) for i in numba.parfor.internal_prange(n): S[i] = hpat.pd_timestamp_ext.parse_datetime_str(str_arr[i]) ret = S f_ir = compile_to_numba_ir( f, { 'hpat': hpat, 'numba': numba }, self.typingctx, (if_series_to_array_type(self.typemap[data.name]), ), self.typemap, self.calltypes) topo_order = find_topo_order(f_ir.blocks) f_ir.blocks[topo_order[-1]].body[-4].target = lhs replace_arg_nodes(f_ir.blocks[topo_order[0]], [data]) return f_ir.blocks
def run(self): dprint_func_ir(self.func_ir, "starting hiframes") topo_order = find_topo_order(self.func_ir.blocks) for label in topo_order: self._get_reverse_copies(self.func_ir.blocks[label].body) new_body = [] for inst in self.func_ir.blocks[label].body: # df['col'] = arr if isinstance( inst, ir.StaticSetItem) and inst.target.name in self.df_vars: df_name = inst.target.name self.df_vars[df_name][inst.index] = inst.value self._update_df_cols() elif isinstance(inst, ir.Assign): out_nodes = self._run_assign(inst) if isinstance(out_nodes, list): new_body.extend(out_nodes) if isinstance(out_nodes, dict): label = include_new_blocks(self.func_ir.blocks, out_nodes, label, new_body) new_body = [] else: new_body.append(inst) self.func_ir.blocks[label].body = new_body self.func_ir._definitions = get_definitions(self.func_ir.blocks) self.func_ir.df_cols = self.df_cols #remove_dead(self.func_ir.blocks, self.func_ir.arg_names) dprint_func_ir(self.func_ir, "after hiframes") if numba.config.DEBUG_ARRAY_OPT == 1: print("df_vars: ", self.df_vars) return
def run(self): dprint_func_ir(self.func_ir, "starting IO") topo_order = find_topo_order(self.func_ir.blocks) for label in topo_order: new_body = [] # copies are collected before running the pass since # variables typed in locals are assigned late self._get_reverse_copies(self.func_ir.blocks[label].body) for inst in self.func_ir.blocks[label].body: if isinstance(inst, ir.Assign): inst_list = self._run_assign(inst) new_body.extend(inst_list) elif isinstance(inst, ir.StaticSetItem): inst_list = self._run_static_setitem(inst) new_body.extend(inst_list) else: new_body.append(inst) self.func_ir.blocks[label].body = new_body # iterative remove dead to make sure all extra code (e.g. df vars) is removed while remove_dead(self.func_ir.blocks, self.func_ir.arg_names, self.func_ir): pass self.func_ir._definitions = get_definitions(self.func_ir.blocks) dprint_func_ir(self.func_ir, "after IO") if debug_prints(): print("h5 files: ", self.h5_files) print("h5 dsets: ", self.h5_dsets)
def _gen_col_var(self, out_var, args, col_var): loc = out_var.loc scope = out_var.scope # calculate mean first mean_var = ir.Var(scope, mk_unique_var("mean_val"), loc) f_mean_blocks = self._gen_col_mean(mean_var, args, col_var) f_mean_blocks = add_offset_to_labels(f_mean_blocks, ir_utils._max_label+1) ir_utils._max_label = max(f_mean_blocks.keys()) m_last_label = find_topo_order(f_mean_blocks)[-1] remove_none_return_from_block(f_mean_blocks[m_last_label]) def f(A, s, m): count = 0 for i in numba.parfor.prange(len(A)): val = A[i] if not np.isnan(val): s += (val-m)**2 count += 1 if count <= 1: s = np.nan else: s = s/(count-1) f_blocks = get_inner_ir(f) replace_var_names(f_blocks, {'A': col_var.name}) replace_var_names(f_blocks, {'s': out_var.name}) replace_var_names(f_blocks, {'m': mean_var.name}) f_blocks[0].body.insert(0, ir.Assign(ir.Const(0.0, loc), out_var, loc)) # attach first var block to last mean block f_mean_blocks[m_last_label].body.extend(f_blocks[0].body) f_blocks.pop(0) f_blocks = add_offset_to_labels(f_blocks, ir_utils._max_label+1) # add offset to jump of first f_block since it didn't go through call f_mean_blocks[m_last_label].body[-1].target += ir_utils._max_label+1 ir_utils._max_label = max(f_blocks.keys()) f_mean_blocks.update(f_blocks) return f_mean_blocks
def run(self): blocks = self.func_ir.blocks array_dists = {} parfor_dists = {} topo_order = find_topo_order(blocks) self._run_analysis(self.func_ir.blocks, topo_order, array_dists, parfor_dists) self.second_pass = True self._run_analysis(self.func_ir.blocks, topo_order, array_dists, parfor_dists) return _dist_analysis_result(array_dists=array_dists, parfor_dists=parfor_dists)
def run(self): dprint_func_ir(self.func_ir, "starting hiframes") topo_order = find_topo_order(self.func_ir.blocks) for label in topo_order: new_body = [] for inst in self.func_ir.blocks[label].body: # df['col'] = arr if isinstance(inst, ir.StaticSetItem) and inst.target.name in self.df_vars: df_name = inst.target.name self.df_vars[df_name][inst.index] = inst.value self._update_df_cols() elif isinstance(inst, ir.Assign): out_nodes = self._run_assign(inst) if isinstance(out_nodes, list): new_body.extend(out_nodes) if isinstance(out_nodes, dict): inner_blocks = add_offset_to_labels(out_nodes, ir_utils._max_label+1) self.func_ir.blocks.update(inner_blocks) ir_utils._max_label = max(self.func_ir.blocks.keys()) scope = self.func_ir.blocks[label].scope loc = self.func_ir.blocks[label].loc inner_topo_order = find_topo_order(inner_blocks) inner_first_label = inner_topo_order[0] inner_last_label = inner_topo_order[-1] remove_none_return_from_block(inner_blocks[inner_last_label]) new_body.append(ir.Jump(inner_first_label, loc)) self.func_ir.blocks[label].body = new_body label = ir_utils.next_label() self.func_ir.blocks[label] = ir.Block(scope, loc) inner_blocks[inner_last_label].body.append(ir.Jump(label, loc)) new_body = [] else: new_body.append(inst) self.func_ir.blocks[label].body = new_body remove_dead(self.func_ir.blocks, self.func_ir.arg_names) dprint_func_ir(self.func_ir, "after hiframes") if config.DEBUG_ARRAY_OPT==1: print("df_vars: ", self.df_vars) return
def run(self): blocks = self.func_ir.blocks array_dists = {} parfor_dists = {} topo_order = find_topo_order(blocks) save_array_dists = {} save_parfor_dists = {1: 1} # dummy value # fixed-point iteration while array_dists != save_array_dists or parfor_dists != save_parfor_dists: save_array_dists = copy.copy(array_dists) save_parfor_dists = copy.copy(parfor_dists) for label in topo_order: self._analyze_block(blocks[label], array_dists, parfor_dists) return _dist_analysis_result(array_dists=array_dists, parfor_dists=parfor_dists)
def include_new_blocks(blocks, new_blocks, label, new_body): inner_blocks = add_offset_to_labels(new_blocks, ir_utils._max_label + 1) blocks.update(inner_blocks) ir_utils._max_label = max(blocks.keys()) scope = blocks[label].scope loc = blocks[label].loc inner_topo_order = find_topo_order(inner_blocks) inner_first_label = inner_topo_order[0] inner_last_label = inner_topo_order[-1] remove_none_return_from_block(inner_blocks[inner_last_label]) new_body.append(ir.Jump(inner_first_label, loc)) blocks[label].body = new_body label = ir_utils.next_label() blocks[label] = ir.Block(scope, loc) inner_blocks[inner_last_label].body.append(ir.Jump(label, loc)) #new_body.clear() return label
def inline_new_blocks(func_ir, block, i, callee_blocks, work_list=None): # adopted from inline_closure_call scope = block.scope instr = block.body[i] # 1. relabel callee_ir by adding an offset callee_blocks = add_offset_to_labels(callee_blocks, ir_utils._max_label + 1) callee_blocks = ir_utils.simplify_CFG(callee_blocks) max_label = max(callee_blocks.keys()) # reset globals in ir_utils before we use it ir_utils._max_label = max_label topo_order = find_topo_order(callee_blocks) # 5. split caller blocks into two new_blocks = [] new_block = ir.Block(scope, block.loc) new_block.body = block.body[i + 1:] new_label = ir_utils.next_label() func_ir.blocks[new_label] = new_block new_blocks.append((new_label, new_block)) block.body = block.body[:i] min_label = topo_order[0] block.body.append(ir.Jump(min_label, instr.loc)) # 6. replace Return with assignment to LHS numba.inline_closurecall._replace_returns(callee_blocks, instr.target, new_label) # remove the old definition of instr.target too if (instr.target.name in func_ir._definitions): func_ir._definitions[instr.target.name] = [] # 7. insert all new blocks, and add back definitions for label in topo_order: # block scope must point to parent's block = callee_blocks[label] block.scope = scope numba.inline_closurecall._add_definitions(func_ir, block) func_ir.blocks[label] = block new_blocks.append((label, block)) if work_list is not None: for block in new_blocks: work_list.append(block) return callee_blocks
def run(self): """run array shape analysis on the IR and save information in array_shape_classes, class_sizes, and array_size_vars (see __init__ comments). May generate some array shape calls if necessary. """ dprint_func_ir(self.func_ir, "starting array analysis") if config.DEBUG_ARRAY_OPT == 1: print("variable types: ", sorted(self.typemap.items())) print("call types: ", self.calltypes) topo_order = find_topo_order(self.func_ir.blocks) for label in topo_order: self._analyze_block(self.func_ir.blocks[label]) self._merge_equivalent_classes() self._cleanup_analysis_data() if config.DEBUG_ARRAY_OPT == 1: self.dump()
def run(self): self._init_run() blocks = self.func_ir.blocks array_dists = {} parfor_dists = {} topo_order = find_topo_order(blocks) self._run_analysis(self.func_ir.blocks, topo_order, array_dists, parfor_dists) self.second_pass = True self._run_analysis(self.func_ir.blocks, topo_order, array_dists, parfor_dists) # rebalance arrays if necessary if auto_rebalance and Distribution.OneD_Var in array_dists.values(): changed = self._rebalance_arrs(array_dists, parfor_dists) if changed: return self.run() return _dist_analysis_result(array_dists=array_dists, parfor_dists=parfor_dists)
def _gen_col_std(self, out_var, args, col_var): loc = out_var.loc scope = out_var.scope # calculate var() first var_var = ir.Var(scope, mk_unique_var("var_val"), loc) f_blocks = self._gen_col_var(var_var, args, col_var) last_label = find_topo_order(f_blocks)[-1] def f(a): a ** 0.5 s_blocks = get_inner_ir(f) replace_var_names(s_blocks, {'a': var_var.name}) remove_none_return_from_block(s_blocks[0]) assert len(s_blocks[0].body) == 2 const_node = s_blocks[0].body[0] pow_node = s_blocks[0].body[1] pow_node.target = out_var f_blocks[last_label].body.insert(-3, const_node) f_blocks[last_label].body.insert(-3, pow_node) return f_blocks
def run(self): dprint_func_ir(self.func_ir, "starting hiframes") topo_order = find_topo_order(self.func_ir.blocks) for label in topo_order: new_body = [] for inst in self.func_ir.blocks[label].body: # df['col'] = arr if isinstance( inst, ir.StaticSetItem) and inst.target.name in self.df_vars: df_name = inst.target.name self.df_vars[df_name][inst.index] = inst.value self._update_df_cols() elif isinstance(inst, ir.Assign): out_nodes = self._run_assign(inst) if isinstance(out_nodes, list): new_body.extend(out_nodes) if isinstance(out_nodes, dict): label = include_new_blocks(self.func_ir.blocks, out_nodes, label, new_body) new_body = [] else: new_body.append(inst) self.func_ir.blocks[label].body = new_body self.func_ir._definitions = _get_definitions(self.func_ir.blocks) #remove_dead(self.func_ir.blocks, self.func_ir.arg_names) if config._has_h5py: io_pass = pio.PIO(self.func_ir, self.locals) io_pass.run() remove_dead(self.func_ir.blocks, self.func_ir.arg_names) DummyFlags = namedtuple('DummyFlags', 'auto_parallel') inline_pass = InlineClosureCallPass(self.func_ir, DummyFlags(True)) inline_pass.run() self.typemap, self.return_type, self.calltypes = numba_compiler.type_inference_stage( self.typingctx, self.func_ir, self.args, None) self.fix_series_filter(self.func_ir.blocks) self.func_ir._definitions = _get_definitions(self.func_ir.blocks) dprint_func_ir(self.func_ir, "after hiframes") if numba.config.DEBUG_ARRAY_OPT == 1: print("df_vars: ", self.df_vars) return
def fix_series_filter(self, blocks): topo_order = find_topo_order(blocks) for label in topo_order: new_body = [] for stmt in blocks[label].body: # find df['col2'] = df['col1'][arr] if (isinstance(stmt, ir.Assign) and isinstance(stmt.value, ir.Expr) and stmt.value.op == 'getitem' and stmt.value.value.name in self.df_cols and stmt.target.name in self.df_cols and self.is_bool_arr(stmt.value.index.name)): lhs = stmt.target in_arr = stmt.value.value index_var = stmt.value.index def f(A, B, ind): for i in numba.parfor.prange(len(A)): s = 0 if ind[i]: s = B[i] else: s = np.nan A[i] = s f_blocks = get_inner_ir(f) replace_var_names(f_blocks, {'A': lhs.name}) replace_var_names(f_blocks, {'B': in_arr.name}) replace_var_names(f_blocks, {'ind': index_var.name}) alloc_nodes = gen_empty_like(in_arr, lhs) f_blocks[0].body = alloc_nodes + f_blocks[0].body label = include_new_blocks(blocks, f_blocks, label, new_body) new_body = [] else: new_body.append(stmt) blocks[label].body = new_body return
def run(self): dprint_func_ir(self.func_ir, "starting IO") topo_order = find_topo_order(self.func_ir.blocks) for label in topo_order: new_body = [] # copies are collected before running the pass since # variables typed in locals are assigned late self._get_reverse_copies(self.func_ir.blocks[label].body) for inst in self.func_ir.blocks[label].body: if isinstance(inst, ir.Assign): inst_list = self._run_assign(inst) new_body.extend(inst_list) elif isinstance(inst, ir.StaticSetItem): inst_list = self._run_static_setitem(inst) new_body.extend(inst_list) else: new_body.append(inst) self.func_ir.blocks[label].body = new_body #remove_dead(self.func_ir.blocks, self.func_ir.arg_names) dprint_func_ir(self.func_ir, "after IO") if config.DEBUG_ARRAY_OPT == 1: print("h5 files: ", self.h5_files) print("h5 dsets: ", self.h5_dsets)
def get_inner_ir(func): # get untyped numba ir f_ir = numba_compiler.run_frontend(func) blocks = f_ir.blocks remove_dels(blocks) topo_order = find_topo_order(blocks) first_block = blocks[topo_order[0]] last_block = blocks[topo_order[-1]] # remove arg nodes new_first_body = [] for stmt in first_block.body: if isinstance(stmt, ir.Assign) and isinstance(stmt.value, ir.Arg): continue new_first_body.append(stmt) first_block.body = new_first_body # rename all variables to avoid conflict, except args var_table = get_name_var_table(blocks) new_var_dict = {} for name, var in var_table.items(): if not (name in f_ir.arg_names): new_var_dict[name] = mk_unique_var(name) replace_var_names(blocks, new_var_dict) return blocks
def _handle_df_col_calls(self, lhs_name, rhs, assign): if guard(find_callname, self.func_ir, rhs) == ('count', 'hpat.hiframes_api'): in_arr = rhs.args[0] f_blocks = compile_to_numba_ir(_column_count_impl, { 'numba': numba, 'np': np, 'hpat': hpat }, self.typingctx, (self.typemap[in_arr.name], ), self.typemap, self.calltypes).blocks topo_order = find_topo_order(f_blocks) first_block = topo_order[0] last_block = topo_order[-1] replace_arg_nodes(f_blocks[first_block], [in_arr]) # assign results to lhs output f_blocks[last_block].body[-4].target = assign.target return f_blocks if guard(find_callname, self.func_ir, rhs) == ('fillna', 'hpat.hiframes_api'): out_arr = rhs.args[0] in_arr = rhs.args[1] val = rhs.args[2] f_blocks = compile_to_numba_ir( _column_fillna_impl, { 'numba': numba, 'np': np }, self.typingctx, (self.typemap[out_arr.name], self.typemap[in_arr.name], self.typemap[val.name]), self.typemap, self.calltypes).blocks first_block = min(f_blocks.keys()) replace_arg_nodes(f_blocks[first_block], [out_arr, in_arr, val]) return f_blocks if guard(find_callname, self.func_ir, rhs) == ('column_sum', 'hpat.hiframes_api'): in_arr = rhs.args[0] f_blocks = compile_to_numba_ir(_column_sum_impl, { 'numba': numba, 'np': np, 'hpat': hpat }, self.typingctx, (self.typemap[in_arr.name], ), self.typemap, self.calltypes).blocks topo_order = find_topo_order(f_blocks) first_block = topo_order[0] last_block = topo_order[-1] replace_arg_nodes(f_blocks[first_block], [in_arr]) # assign results to lhs output f_blocks[last_block].body[-4].target = assign.target return f_blocks if guard(find_callname, self.func_ir, rhs) == ('mean', 'hpat.hiframes_api'): in_arr = rhs.args[0] f_blocks = compile_to_numba_ir(_column_mean_impl, { 'numba': numba, 'np': np, 'hpat': hpat }, self.typingctx, (self.typemap[in_arr.name], ), self.typemap, self.calltypes).blocks topo_order = find_topo_order(f_blocks) first_block = topo_order[0] last_block = topo_order[-1] replace_arg_nodes(f_blocks[first_block], [in_arr]) # assign results to lhs output f_blocks[last_block].body[-4].target = assign.target return f_blocks if guard(find_callname, self.func_ir, rhs) == ('var', 'hpat.hiframes_api'): in_arr = rhs.args[0] f_blocks = compile_to_numba_ir(_column_var_impl, { 'numba': numba, 'np': np, 'hpat': hpat }, self.typingctx, (self.typemap[in_arr.name], ), self.typemap, self.calltypes).blocks topo_order = find_topo_order(f_blocks) first_block = topo_order[0] last_block = topo_order[-1] replace_arg_nodes(f_blocks[first_block], [in_arr]) # assign results to lhs output f_blocks[last_block].body[-4].target = assign.target return f_blocks return
def _fix_nested_array(func_ir): """Look for assignment like: a[..] = b, where both a and b are numpy arrays, and try to eliminate array b by expanding a with an extra dimension. """ blocks = func_ir.blocks cfg = compute_cfg_from_blocks(blocks) usedefs = compute_use_defs(blocks) empty_deadmap = dict([(label, set()) for label in blocks.keys()]) livemap = compute_live_variables(cfg, blocks, usedefs.defmap, empty_deadmap) def find_array_def(arr): """Find numpy array definition such as arr = numba.unsafe.ndarray.empty_inferred(...). If it is arr = b[...], find array definition of b recursively. """ arr_def = func_ir.get_definition(arr) _make_debug_print("find_array_def")(arr, arr_def) if isinstance(arr_def, ir.Expr): if guard(_find_unsafe_empty_inferred, func_ir, arr_def): return arr_def elif arr_def.op == 'getitem': return find_array_def(arr_def.value) raise GuardException def fix_dependencies(expr, varlist): """Double check if all variables in varlist are defined before expr is used. Try to move constant definition when the check fails. Bails out by raising GuardException if it can't be moved. """ debug_print = _make_debug_print("fix_dependencies") for label, block in blocks.items(): scope = block.scope body = block.body defined = set() for i in range(len(body)): inst = body[i] if isinstance(inst, ir.Assign): defined.add(inst.target.name) if inst.value == expr: new_varlist = [] for var in varlist: # var must be defined before this inst, or live # and not later defined. if (var.name in defined or (var.name in livemap[label] and not (var.name in usedefs.defmap[label]))): debug_print(var.name, " already defined") new_varlist.append(var) else: debug_print(var.name, " not yet defined") var_def = get_definition(func_ir, var.name) if isinstance(var_def, ir.Const): loc = var.loc new_var = ir.Var(scope, mk_unique_var("new_var"), loc) new_const = ir.Const(var_def.value, loc) new_vardef = _new_definition(func_ir, new_var, new_const, loc) new_body = [] new_body.extend(body[:i]) new_body.append(new_vardef) new_body.extend(body[i:]) block.body = new_body new_varlist.append(new_var) else: raise GuardException return new_varlist # when expr is not found in block raise GuardException def fix_array_assign(stmt): """For assignment like lhs[idx] = rhs, where both lhs and rhs are arrays, do the following: 1. find the definition of rhs, which has to be a call to numba.unsafe.ndarray.empty_inferred 2. find the source array creation for lhs, insert an extra dimension of size of b. 3. replace the definition of rhs = numba.unsafe.ndarray.empty_inferred(...) with rhs = lhs[idx] """ require(isinstance(stmt, ir.SetItem)) require(isinstance(stmt.value, ir.Var)) debug_print = _make_debug_print("fix_array_assign") debug_print("found SetItem: ", stmt) lhs = stmt.target # Find the source array creation of lhs lhs_def = find_array_def(lhs) debug_print("found lhs_def: ", lhs_def) rhs_def = get_definition(func_ir, stmt.value) debug_print("found rhs_def: ", rhs_def) require(isinstance(rhs_def, ir.Expr)) if rhs_def.op == 'cast': rhs_def = get_definition(func_ir, rhs_def.value) require(isinstance(rhs_def, ir.Expr)) require(_find_unsafe_empty_inferred(func_ir, rhs_def)) # Find the array dimension of rhs dim_def = get_definition(func_ir, rhs_def.args[0]) require(isinstance(dim_def, ir.Expr) and dim_def.op == 'build_tuple') debug_print("dim_def = ", dim_def) extra_dims = [ get_definition(func_ir, x, lhs_only=True) for x in dim_def.items ] debug_print("extra_dims = ", extra_dims) # Expand size tuple when creating lhs_def with extra_dims size_tuple_def = get_definition(func_ir, lhs_def.args[0]) require(isinstance(size_tuple_def, ir.Expr) and size_tuple_def.op == 'build_tuple') debug_print("size_tuple_def = ", size_tuple_def) extra_dims = fix_dependencies(size_tuple_def, extra_dims) size_tuple_def.items += extra_dims # In-place modify rhs_def to be getitem rhs_def.op = 'getitem' rhs_def.value = get_definition(func_ir, lhs, lhs_only=True) rhs_def.index = stmt.index del rhs_def._kws['func'] del rhs_def._kws['args'] del rhs_def._kws['vararg'] del rhs_def._kws['kws'] # success return True for label in find_topo_order(func_ir.blocks): block = func_ir.blocks[label] for stmt in block.body: if guard(fix_array_assign, stmt): block.body.remove(stmt)
def inline_closure_call(func_ir, glbls, block, i, callee, typingctx=None, arg_typs=None, typemap=None, calltypes=None, work_list=None): """Inline the body of `callee` at its callsite (`i`-th instruction of `block`) `func_ir` is the func_ir object of the caller function and `glbls` is its global variable environment (func_ir.func_id.func.__globals__). `block` is the IR block of the callsite and `i` is the index of the callsite's node. `callee` is either the called function or a make_function node. `typingctx`, `typemap` and `calltypes` are typing data structures of the caller, available if we are in a typed pass. `arg_typs` includes the types of the arguments at the callsite. """ scope = block.scope instr = block.body[i] call_expr = instr.value debug_print = _make_debug_print("inline_closure_call") debug_print("Found closure call: ", instr, " with callee = ", callee) # support both function object and make_function Expr callee_code = callee.code if hasattr(callee, 'code') else callee.__code__ callee_defaults = callee.defaults if hasattr(callee, 'defaults') else callee.__defaults__ callee_closure = callee.closure if hasattr(callee, 'closure') else callee.__closure__ # first, get the IR of the callee callee_ir = get_ir_of_code(glbls, callee_code) callee_blocks = callee_ir.blocks # 1. relabel callee_ir by adding an offset max_label = max(func_ir.blocks.keys()) callee_blocks = add_offset_to_labels(callee_blocks, max_label + 1) callee_blocks = simplify_CFG(callee_blocks) callee_ir.blocks = callee_blocks min_label = min(callee_blocks.keys()) max_label = max(callee_blocks.keys()) # reset globals in ir_utils before we use it ir_utils._max_label = max_label debug_print("After relabel") _debug_dump(callee_ir) # 2. rename all local variables in callee_ir with new locals created in func_ir callee_scopes = _get_all_scopes(callee_blocks) debug_print("callee_scopes = ", callee_scopes) # one function should only have one local scope assert(len(callee_scopes) == 1) callee_scope = callee_scopes[0] var_dict = {} for var in callee_scope.localvars._con.values(): if not (var.name in callee_code.co_freevars): new_var = scope.define(mk_unique_var(var.name), loc=var.loc) var_dict[var.name] = new_var debug_print("var_dict = ", var_dict) replace_vars(callee_blocks, var_dict) debug_print("After local var rename") _debug_dump(callee_ir) # 3. replace formal parameters with actual arguments args = list(call_expr.args) if callee_defaults: debug_print("defaults = ", callee_defaults) if isinstance(callee_defaults, tuple): # Python 3.5 args = args + list(callee_defaults) elif isinstance(callee_defaults, ir.Var) or isinstance(callee_defaults, str): defaults = func_ir.get_definition(callee_defaults) assert(isinstance(defaults, ir.Const)) loc = defaults.loc args = args + [ir.Const(value=v, loc=loc) for v in defaults.value] else: raise NotImplementedError( "Unsupported defaults to make_function: {}".format(defaults)) debug_print("After arguments rename: ") _debug_dump(callee_ir) # 4. replace freevar with actual closure var if callee_closure: closure = func_ir.get_definition(callee_closure) debug_print("callee's closure = ", closure) if isinstance(closure, tuple): cellget = ctypes.pythonapi.PyCell_Get cellget.restype = ctypes.py_object cellget.argtypes = (ctypes.py_object,) items = tuple(cellget(x) for x in closure) else: assert(isinstance(closure, ir.Expr) and closure.op == 'build_tuple') items = closure.items assert(len(callee_code.co_freevars) == len(items)) _replace_freevars(callee_blocks, items) debug_print("After closure rename") _debug_dump(callee_ir) if typingctx: from numba import compiler f_typemap, f_return_type, f_calltypes = compiler.type_inference_stage( typingctx, callee_ir, arg_typs, None) canonicalize_array_math(callee_ir, f_typemap, f_calltypes, typingctx) # remove argument entries like arg.a from typemap arg_names = [vname for vname in f_typemap if vname.startswith("arg.")] for a in arg_names: f_typemap.pop(a) typemap.update(f_typemap) calltypes.update(f_calltypes) _replace_args_with(callee_blocks, args) # 5. split caller blocks into two new_blocks = [] new_block = ir.Block(scope, block.loc) new_block.body = block.body[i + 1:] new_label = next_label() func_ir.blocks[new_label] = new_block new_blocks.append((new_label, new_block)) block.body = block.body[:i] block.body.append(ir.Jump(min_label, instr.loc)) # 6. replace Return with assignment to LHS topo_order = find_topo_order(callee_blocks) _replace_returns(callee_blocks, instr.target, new_label) # remove the old definition of instr.target too if (instr.target.name in func_ir._definitions): func_ir._definitions[instr.target.name] = [] # 7. insert all new blocks, and add back definitions for label in topo_order: # block scope must point to parent's block = callee_blocks[label] block.scope = scope _add_definitions(func_ir, block) func_ir.blocks[label] = block new_blocks.append((label, block)) debug_print("After merge in") _debug_dump(func_ir) if work_list != None: for block in new_blocks: work_list.append(block) return callee_blocks
def inline_closure_call(self, block, i, callee): """Inline the body of `callee` at its callsite (`i`-th instruction of `block`) """ scope = block.scope instr = block.body[i] call_expr = instr.value debug_print = _make_debug_print("inline_closure_call") debug_print("Found closure call: ", instr, " with callee = ", callee) func_ir = self.func_ir # first, get the IR of the callee callee_ir = self.get_ir_of_code(callee.code) callee_blocks = callee_ir.blocks # 1. relabel callee_ir by adding an offset max_label = max(func_ir.blocks.keys()) callee_blocks = add_offset_to_labels(callee_blocks, max_label + 1) callee_ir.blocks = callee_blocks min_label = min(callee_blocks.keys()) max_label = max(callee_blocks.keys()) # reset globals in ir_utils before we use it ir_utils._max_label = max_label debug_print("After relabel") _debug_dump(callee_ir) # 2. rename all local variables in callee_ir with new locals created in func_ir callee_scopes = _get_all_scopes(callee_blocks) debug_print("callee_scopes = ", callee_scopes) # one function should only have one local scope assert(len(callee_scopes) == 1) callee_scope = callee_scopes[0] var_dict = {} for var in callee_scope.localvars._con.values(): if not (var.name in callee.code.co_freevars): new_var = scope.define(mk_unique_var(var.name), loc=var.loc) var_dict[var.name] = new_var debug_print("var_dict = ", var_dict) replace_vars(callee_blocks, var_dict) debug_print("After local var rename") _debug_dump(callee_ir) # 3. replace formal parameters with actual arguments args = list(call_expr.args) if callee.defaults: debug_print("defaults = ", callee.defaults) if isinstance(callee.defaults, tuple): # Python 3.5 args = args + list(callee.defaults) elif isinstance(callee.defaults, ir.Var) or isinstance(callee.defaults, str): defaults = func_ir.get_definition(callee.defaults) assert(isinstance(defaults, ir.Const)) loc = defaults.loc args = args + [ir.Const(value=v, loc=loc) for v in defaults.value] else: raise NotImplementedError( "Unsupported defaults to make_function: {}".format(defaults)) _replace_args_with(callee_blocks, args) debug_print("After arguments rename: ") _debug_dump(callee_ir) # 4. replace freevar with actual closure var if callee.closure: closure = func_ir.get_definition(callee.closure) assert(isinstance(closure, ir.Expr) and closure.op == 'build_tuple') assert(len(callee.code.co_freevars) == len(closure.items)) debug_print("callee's closure = ", closure) _replace_freevars(callee_blocks, closure.items) debug_print("After closure rename") _debug_dump(callee_ir) # 5. split caller blocks into two new_blocks = [] new_block = ir.Block(scope, block.loc) new_block.body = block.body[i + 1:] new_label = next_label() func_ir.blocks[new_label] = new_block new_blocks.append((new_label, new_block)) block.body = block.body[:i] block.body.append(ir.Jump(min_label, instr.loc)) # 6. replace Return with assignment to LHS topo_order = find_topo_order(callee_blocks) _replace_returns(callee_blocks, instr.target, new_label) # remove the old definition of instr.target too if (instr.target.name in func_ir._definitions): func_ir._definitions[instr.target.name] = [] # 7. insert all new blocks, and add back definitions for label in topo_order: # block scope must point to parent's block = callee_blocks[label] block.scope = scope _add_definitions(func_ir, block) func_ir.blocks[label] = block new_blocks.append((label, block)) debug_print("After merge in") _debug_dump(func_ir) return new_blocks
def _fix_nested_array(func_ir): """Look for assignment like: a[..] = b, where both a and b are numpy arrays, and try to eliminate array b by expanding a with an extra dimension. """ """ cfg = compute_cfg_from_blocks(func_ir.blocks) all_loops = list(cfg.loops().values()) def find_nest_level(label): level = 0 for loop in all_loops: if label in loop.body: level += 1 """ def find_array_def(arr): """Find numpy array definition such as arr = numba.unsafe.ndarray.empty_inferred(...). If it is arr = b[...], find array definition of b recursively. """ arr_def = func_ir.get_definition(arr) _make_debug_print("find_array_def")(arr, arr_def) if isinstance(arr_def, ir.Expr): if guard(_find_unsafe_empty_inferred, func_ir, arr_def): return arr_def elif arr_def.op == 'getitem': return find_array_def(arr_def.value) raise GuardException def fix_array_assign(stmt): """For assignment like lhs[idx] = rhs, where both lhs and rhs are arrays, do the following: 1. find the definition of rhs, which has to be a call to numba.unsafe.ndarray.empty_inferred 2. find the source array creation for lhs, insert an extra dimension of size of b. 3. replace the definition of rhs = numba.unsafe.ndarray.empty_inferred(...) with rhs = lhs[idx] """ require(isinstance(stmt, ir.SetItem)) require(isinstance(stmt.value, ir.Var)) debug_print = _make_debug_print("fix_array_assign") debug_print("found SetItem: ", stmt) lhs = stmt.target # Find the source array creation of lhs lhs_def = find_array_def(lhs) debug_print("found lhs_def: ", lhs_def) rhs_def = get_definition(func_ir, stmt.value) debug_print("found rhs_def: ", rhs_def) require(isinstance(rhs_def, ir.Expr)) if rhs_def.op == 'cast': rhs_def = get_definition(func_ir, rhs_def.value) require(isinstance(rhs_def, ir.Expr)) require(_find_unsafe_empty_inferred(func_ir, rhs_def)) # Find the array dimension of rhs dim_def = get_definition(func_ir, rhs_def.args[0]) require(isinstance(dim_def, ir.Expr) and dim_def.op == 'build_tuple') debug_print("dim_def = ", dim_def) extra_dims = [ get_definition(func_ir, x, lhs_only=True) for x in dim_def.items ] debug_print("extra_dims = ", extra_dims) # Expand size tuple when creating lhs_def with extra_dims size_tuple_def = get_definition(func_ir, lhs_def.args[0]) require(isinstance(size_tuple_def, ir.Expr) and size_tuple_def.op == 'build_tuple') debug_print("size_tuple_def = ", size_tuple_def) size_tuple_def.items += extra_dims # In-place modify rhs_def to be getitem rhs_def.op = 'getitem' rhs_def.value = get_definition(func_ir, lhs, lhs_only=True) rhs_def.index = stmt.index del rhs_def._kws['func'] del rhs_def._kws['args'] del rhs_def._kws['vararg'] del rhs_def._kws['kws'] # success return True for label in find_topo_order(func_ir.blocks): block = func_ir.blocks[label] for stmt in block.body: if guard(fix_array_assign, stmt): block.body.remove(stmt)
def run(self): blocks = self.func_ir.blocks topo_order = find_topo_order(blocks) for label in topo_order: new_body = [] for inst in blocks[label].body: if isinstance(inst, ir.Assign): out_nodes = self._run_assign(inst) if isinstance(out_nodes, list): new_body.extend(out_nodes) if isinstance(out_nodes, dict): label = include_new_blocks(blocks, out_nodes, label, new_body) new_body = [] if isinstance(out_nodes, tuple): gen_blocks, post_nodes = out_nodes label = include_new_blocks(blocks, gen_blocks, label, new_body) new_body = post_nodes else: new_body.append(inst) blocks[label].body = new_body if debug_prints(): # pragma: no cover print("--- types before Series replacement:", self.typemap) print("calltypes: ", self.calltypes) replace_series = {} for vname, typ in self.typemap.items(): if isinstance(typ, SeriesType): # print("replacing series type", vname) new_typ = series_to_array_type(typ) replace_series[vname] = new_typ # replace array.call() variable types if isinstance(typ, types.BoundFunction) and isinstance( typ.this, SeriesType): this = series_to_array_type(typ.this) # TODO: handle string arrays, etc. assert typ.typing_key.startswith('array.') attr = typ.typing_key[len('array.'):] resolver = getattr(ArrayAttribute, 'resolve_' + attr) # methods are either installed with install_array_method or # using @bound_function in arraydecl.py if hasattr(resolver, '__wrapped__'): resolver = bound_function(typ.typing_key)( resolver.__wrapped__) new_typ = resolver(ArrayAttribute(self.typingctx), this) replace_series[vname] = new_typ for vname, typ in replace_series.items(): self.typemap.pop(vname) self.typemap[vname] = typ replace_calltype = {} # replace sig of getitem/setitem/... series type with array for call, sig in self.calltypes.items(): if sig is None: continue assert isinstance(sig, Signature) sig.return_type = if_series_to_array_type(sig.return_type) sig.args = tuple(map(if_series_to_array_type, sig.args)) # XXX: side effect: force update of call signatures if isinstance(call, ir.Expr) and call.op == 'call': # StencilFunc requires kws for typing so sig.args can't be used # reusing sig.args since some types become Const in sig argtyps = sig.args[:len(call.args)] kwtyps = {name: self.typemap[v.name] for name, v in call.kws} new_sig = self.typemap[call.func.name].get_call_type( self.typingctx, argtyps, kwtyps) # calltypes of things like BoundFunction (array.call) need to # be update for lowering to work # XXX: new_sig could be None for things like np.int32() if call in self.calltypes and new_sig is not None: old_sig = self.calltypes[call] # fix types with undefined dtypes in empty_inferred, etc. return_type = _fix_typ_undefs(new_sig.return_type, old_sig.return_type) args = tuple( _fix_typ_undefs(a, b) for a, b in zip(new_sig.args, old_sig.args)) replace_calltype[call] = Signature(return_type, args, new_sig.recvr, new_sig.pysig) for call, sig in replace_calltype.items(): self.calltypes.pop(call) self.calltypes[call] = sig if debug_prints(): # pragma: no cover print("--- types after Series replacement:", self.typemap) print("calltypes: ", self.calltypes) self.func_ir._definitions = get_definitions(self.func_ir.blocks) return if_series_to_unbox(self.return_type)
def _handle_df_col_calls(self, assign, lhs, rhs, func_name): if func_name == 'count': in_arr = rhs.args[0] f_blocks = compile_to_numba_ir( _column_count_impl, { 'numba': numba, 'np': np, 'hpat': hpat }, self.typingctx, (if_series_to_array_type(self.typemap[in_arr.name]), ), self.typemap, self.calltypes).blocks topo_order = find_topo_order(f_blocks) first_block = topo_order[0] last_block = topo_order[-1] replace_arg_nodes(f_blocks[first_block], [in_arr]) # assign results to lhs output f_blocks[last_block].body[-3].target = assign.target return f_blocks if func_name == 'fillna': out_arr = rhs.args[0] in_arr = rhs.args[1] val = rhs.args[2] f_blocks = compile_to_numba_ir( _column_fillna_impl, { 'numba': numba, 'np': np }, self.typingctx, (if_series_to_array_type(self.typemap[out_arr.name]), if_series_to_array_type(self.typemap[in_arr.name]), if_series_to_array_type(self.typemap[val.name])), self.typemap, self.calltypes).blocks first_block = min(f_blocks.keys()) replace_arg_nodes(f_blocks[first_block], [out_arr, in_arr, val]) return f_blocks if func_name == 'column_sum': in_arr = rhs.args[0] f_blocks = compile_to_numba_ir( _column_sum_impl, { 'numba': numba, 'np': np, 'hpat': hpat }, self.typingctx, (if_series_to_array_type(self.typemap[in_arr.name]), ), self.typemap, self.calltypes).blocks topo_order = find_topo_order(f_blocks) first_block = topo_order[0] last_block = topo_order[-1] replace_arg_nodes(f_blocks[first_block], [in_arr]) # assign results to lhs output f_blocks[last_block].body[-3].target = assign.target return f_blocks if func_name == 'mean': in_arr = rhs.args[0] f_blocks = compile_to_numba_ir( _column_mean_impl, { 'numba': numba, 'np': np, 'hpat': hpat }, self.typingctx, (if_series_to_array_type(self.typemap[in_arr.name]), ), self.typemap, self.calltypes).blocks topo_order = find_topo_order(f_blocks) first_block = topo_order[0] last_block = topo_order[-1] replace_arg_nodes(f_blocks[first_block], [in_arr]) # assign results to lhs output f_blocks[last_block].body[-3].target = assign.target return f_blocks if func_name == 'var': in_arr = rhs.args[0] f_blocks = compile_to_numba_ir( _column_var_impl, { 'numba': numba, 'np': np, 'hpat': hpat }, self.typingctx, (if_series_to_array_type(self.typemap[in_arr.name]), ), self.typemap, self.calltypes).blocks topo_order = find_topo_order(f_blocks) first_block = topo_order[0] last_block = topo_order[-1] replace_arg_nodes(f_blocks[first_block], [in_arr]) # assign results to lhs output f_blocks[last_block].body[-3].target = assign.target return f_blocks return [assign]
def fixes_after_typing(self, blocks): call_table, _ = ir_utils.get_call_table(self.func_ir.blocks) topo_order = find_topo_order(blocks) for label in topo_order: new_body = [] for stmt in blocks[label].body: # convert str_arr==str into parfor if (isinstance(stmt, ir.Assign) and isinstance(stmt.value, ir.Expr) and stmt.value.op == 'binop' and stmt.value.fn in ['==', '!='] and (self.typemap[stmt.value.lhs.name] == string_array_type or self.typemap[stmt.value.rhs.name] == string_array_type)): lhs = stmt.value.lhs rhs = stmt.value.rhs lhs_access = 'A' rhs_access = 'B' len_call = 'A.size' if self.typemap[lhs.name] == string_array_type: lhs_access = 'A[i]' if self.typemap[rhs.name] == string_array_type: lhs_access = 'B[i]' len_call = 'B.size' func_text = 'def f(A, B):\n' func_text += ' l = {}\n'.format(len_call) func_text += ' S = np.empty(l, dtype=np.bool_)\n' func_text += ' for i in numba.parfor.prange(l):\n' func_text += ' S[i] = {} {} {}\n'.format(lhs_access, stmt.value.fn, rhs_access) loc_vars = {} exec(func_text, {}, loc_vars) f = loc_vars['f'] f_blocks = compile_to_numba_ir(f, {'numba': numba, 'np': np}).blocks replace_arg_nodes(f_blocks[min(f_blocks.keys())], [lhs, rhs]) label = include_new_blocks(blocks, f_blocks, label, new_body) new_body = [] # replace == expression with result of parfor (S) # S is target of last statement in 1st block of f stmt.value = f_blocks[min(f_blocks.keys())].body[-2].target # arr = fix_df_array(col) -> arr=col if col is array if (isinstance(stmt, ir.Assign) and isinstance(stmt.value, ir.Expr) and stmt.value.op == 'call' and stmt.value.func.name in call_table and call_table[stmt.value.func.name] == ['fix_df_array', 'hiframes_api', hpat] and isinstance(self.typemap[stmt.value.args[0].name], (types.Array, StringArrayType))): stmt.value = stmt.value.args[0] # find df['col2'] = df['col1'][arr] if (isinstance(stmt, ir.Assign) and isinstance(stmt.value, ir.Expr) and stmt.value.op=='getitem' and stmt.value.value.name in self.df_cols and stmt.target.name in self.df_cols and self.is_bool_arr(stmt.value.index.name)): lhs = stmt.target in_arr = stmt.value.value index_var = stmt.value.index def f(A, B, ind): for i in numba.parfor.prange(len(A)): s = 0 if ind[i]: s = B[i] else: s= np.nan A[i] = s f_blocks = get_inner_ir(f) replace_var_names(f_blocks, {'A': lhs.name}) replace_var_names(f_blocks, {'B': in_arr.name}) replace_var_names(f_blocks, {'ind': index_var.name}) alloc_nodes = gen_empty_like(in_arr, lhs) f_blocks[0].body = alloc_nodes + f_blocks[0].body label = include_new_blocks(blocks, f_blocks, label, new_body) new_body = [] else: new_body.append(stmt) blocks[label].body = new_body return