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 hoist(parfor_params, loop_body, typemap, wrapped_blocks): dep_on_param = copy.copy(parfor_params) hoisted = [] def_once = compute_def_once(loop_body) (call_table, reverse_call_table) = get_call_table(wrapped_blocks) for label, block in loop_body.items(): new_block = [] for inst in block.body: if isinstance(inst, ir.Assign) and inst.target.name in def_once: if _hoist_internal(inst, dep_on_param, call_table, hoisted, typemap): # don't add this instuction to the block since it is hoisted continue elif isinstance(inst, parfor.Parfor): new_init_block = [] if config.DEBUG_ARRAY_OPT == 1: print("parfor") inst.dump() for ib_inst in inst.init_block.body: if (isinstance(ib_inst, ir.Assign) and ib_inst.target.name in def_once): if _hoist_internal(ib_inst, dep_on_param, call_table, hoisted, typemap): # don't add this instuction to the block since it is hoisted continue new_init_block.append(ib_inst) inst.init_block.body = new_init_block new_block.append(inst) block.body = new_block return hoisted
def run(self): """ Finds all calls to StencilFuncs in the IR and converts them to parfor. """ from numba.stencil import StencilFunc # Get all the calls in the function IR. call_table, _ = get_call_table(self.func_ir.blocks) stencil_calls = [] stencil_dict = {} for call_varname, call_list in call_table.items(): if isinstance(call_list[0], StencilFunc): # Remember all calls to StencilFuncs. stencil_calls.append(call_varname) stencil_dict[call_varname] = call_list[0] if not stencil_calls: return # return early if no stencil calls found # find and transform stencil calls for label, block in self.func_ir.blocks.items(): for i, stmt in reversed(list(enumerate(block.body))): # Found a call to a StencilFunc. if (isinstance(stmt, ir.Assign) and isinstance(stmt.value, ir.Expr) and stmt.value.op == 'call' and stmt.value.func.name in stencil_calls): kws = dict(stmt.value.kws) # Create dictionary of input argument number to # the argument itself. input_dict = { i: stmt.value.args[i] for i in range(len(stmt.value.args)) } in_args = stmt.value.args arg_typemap = tuple(self.typemap[i.name] for i in in_args) for arg_type in arg_typemap: if isinstance(arg_type, types.BaseTuple): raise ValueError("Tuple parameters not supported " \ "for stencil kernels in parallel=True mode.") out_arr = kws.get('out') # Get the StencilFunc object corresponding to this call. sf = stencil_dict[stmt.value.func.name] stencil_blocks, rt, arg_to_arr_dict = get_stencil_blocks( sf, self.typingctx, arg_typemap, block.scope, block.loc, input_dict, self.typemap, self.calltypes) index_offsets = sf.options.get('index_offsets', None) gen_nodes = self._mk_stencil_parfor( label, in_args, out_arr, stencil_blocks, index_offsets, stmt.target, rt, sf, arg_to_arr_dict) block.body = block.body[:i] + gen_nodes + block.body[i + 1:] # Found a call to a stencil via numba.stencil(). elif (isinstance(stmt, ir.Assign) and isinstance(stmt.value, ir.Expr) and stmt.value.op == 'call' and guard(find_callname, self.func_ir, stmt.value) == ('stencil', 'numba')): # remove dummy stencil() call stmt.value = ir.Const(0, stmt.loc)
def hoist(parfor_params, loop_body, typemap, wrapped_blocks): dep_on_param = copy.copy(parfor_params) hoisted = [] def_once = compute_def_once(loop_body) (call_table, reverse_call_table) = get_call_table(wrapped_blocks) for label, block in loop_body.items(): new_block = [] for inst in block.body: if isinstance(inst, ir.Assign) and inst.target.name in def_once: if _hoist_internal(inst, dep_on_param, call_table, hoisted, typemap): # don't add this instuction to the block since it is hoisted continue elif isinstance(inst, parfor.Parfor): new_init_block = [] if config.DEBUG_ARRAY_OPT == 1: print("parfor") inst.dump() for ib_inst in inst.init_block.body: if (isinstance(ib_inst, ir.Assign) and ib_inst.target.name in def_once): if _hoist_internal(ib_inst, dep_on_param, call_table, hoisted, typemap): # don't add this instuction to the block since it is hoisted continue new_init_block.append(ib_inst) inst.init_block.body = new_init_block new_block.append(inst) block.body = new_block return hoisted
def __init__(self, func_ir, typemap, calltypes): self.func_ir = func_ir self.typemap = typemap self.calltypes = calltypes self._call_table, _ = get_call_table(func_ir.blocks) self._tuple_table = get_tuple_table(func_ir.blocks) self._parallel_accesses = set() self._T_arrs = set()
def run(self): """ Finds all calls to StencilFuncs in the IR and converts them to parfor. """ from numba.stencil import StencilFunc # Get all the calls in the function IR. call_table, _ = get_call_table(self.func_ir.blocks) stencil_calls = [] stencil_dict = {} for call_varname, call_list in call_table.items(): if isinstance(call_list[0], StencilFunc): # Remember all calls to StencilFuncs. stencil_calls.append(call_varname) stencil_dict[call_varname] = call_list[0] if not stencil_calls: return # return early if no stencil calls found # find and transform stencil calls for label, block in self.func_ir.blocks.items(): for i, stmt in reversed(list(enumerate(block.body))): # Found a call to a StencilFunc. if (isinstance(stmt, ir.Assign) and isinstance(stmt.value, ir.Expr) and stmt.value.op == 'call' and stmt.value.func.name in stencil_calls): kws = dict(stmt.value.kws) # Create dictionary of input argument number to # the argument itself. input_dict = {i: stmt.value.args[i] for i in range(len(stmt.value.args))} in_args = stmt.value.args arg_typemap = tuple(self.typemap[i.name] for i in in_args) for arg_type in arg_typemap: if isinstance(arg_type, types.BaseTuple): raise ValueError("Tuple parameters not supported " \ "for stencil kernels in parallel=True mode.") out_arr = kws.get('out') # Get the StencilFunc object corresponding to this call. sf = stencil_dict[stmt.value.func.name] stencil_ir, rt, arg_to_arr_dict = get_stencil_ir(sf, self.typingctx, arg_typemap, block.scope, block.loc, input_dict, self.typemap, self.calltypes) index_offsets = sf.options.get('index_offsets', None) gen_nodes = self._mk_stencil_parfor(label, in_args, out_arr, stencil_ir, index_offsets, stmt.target, rt, sf, arg_to_arr_dict) block.body = block.body[:i] + gen_nodes + block.body[i+1:] # Found a call to a stencil via numba.stencil(). elif (isinstance(stmt, ir.Assign) and isinstance(stmt.value, ir.Expr) and stmt.value.op == 'call' and guard(find_callname, self.func_ir, stmt.value) == ('stencil', 'numba')): # remove dummy stencil() call stmt.value = ir.Const(0, stmt.loc)
def inline_calls(func_ir): call_table, _ = ir_utils.get_call_table(func_ir.blocks) for label, block in func_ir.blocks.items(): for i, stmt in enumerate(block.body): if isinstance(stmt, ir.Assign): rhs = stmt.value if isinstance(rhs, ir.Expr) and rhs.op == 'call': func = rhs.func.name if (func in call_table and call_table[func] and isinstance( call_table[func][0], CPUDispatcher)): py_func = call_table[func][0].py_func inline_calls_inner(func_ir, block, stmt, i, py_func) return # inline_calls_inner will call back recursively
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