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.core.inline_closurecall._add_definitions(func_ir, block) work_list.append((_label, block)) 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.core.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.core.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 get_stencil_ir(sf, typingctx, args, scope, loc, input_dict, typemap, calltypes): """get typed IR from stencil bytecode """ from numba.core.cpu import CPUContext from numba.core.registry import cpu_target from numba.core.annotations import type_annotations from numba.core.typed_passes import type_inference_stage # get untyped IR stencil_func_ir = sf.kernel_ir.copy() # copy the IR nodes to avoid changing IR in the StencilFunc object stencil_blocks = copy.deepcopy(stencil_func_ir.blocks) stencil_func_ir.blocks = stencil_blocks name_var_table = ir_utils.get_name_var_table(stencil_func_ir.blocks) if "out" in name_var_table: raise ValueError( "Cannot use the reserved word 'out' in stencil kernels.") # get typed IR with a dummy pipeline (similar to test_parfors.py) targetctx = CPUContext(typingctx) with cpu_target.nested_context(typingctx, targetctx): tp = DummyPipeline(typingctx, targetctx, args, stencil_func_ir) rewrites.rewrite_registry.apply('before-inference', tp.state) tp.state.typemap, tp.state.return_type, tp.state.calltypes = type_inference_stage( tp.state.typingctx, tp.state.func_ir, tp.state.args, None) type_annotations.TypeAnnotation(func_ir=tp.state.func_ir, typemap=tp.state.typemap, calltypes=tp.state.calltypes, lifted=(), lifted_from=None, args=tp.state.args, return_type=tp.state.return_type, html_output=config.HTML) # make block labels unique stencil_blocks = ir_utils.add_offset_to_labels(stencil_blocks, ir_utils.next_label()) min_label = min(stencil_blocks.keys()) max_label = max(stencil_blocks.keys()) ir_utils._max_label = max_label if config.DEBUG_ARRAY_OPT >= 1: print("Initial stencil_blocks") ir_utils.dump_blocks(stencil_blocks) # rename variables, var_dict = {} for v, typ in tp.state.typemap.items(): new_var = ir.Var(scope, mk_unique_var(v), loc) var_dict[v] = new_var typemap[new_var.name] = typ # add new var type for overall function ir_utils.replace_vars(stencil_blocks, var_dict) if config.DEBUG_ARRAY_OPT >= 1: print("After replace_vars") ir_utils.dump_blocks(stencil_blocks) # add call types to overall function for call, call_typ in tp.state.calltypes.items(): calltypes[call] = call_typ arg_to_arr_dict = {} # replace arg with arr for block in stencil_blocks.values(): for stmt in block.body: if isinstance(stmt, ir.Assign) and isinstance(stmt.value, ir.Arg): if config.DEBUG_ARRAY_OPT >= 1: print("input_dict", input_dict, stmt.value.index, stmt.value.name, stmt.value.index in input_dict) arg_to_arr_dict[stmt.value.name] = input_dict[ stmt.value.index].name stmt.value = input_dict[stmt.value.index] if config.DEBUG_ARRAY_OPT >= 1: print("arg_to_arr_dict", arg_to_arr_dict) print("After replace arg with arr") ir_utils.dump_blocks(stencil_blocks) ir_utils.remove_dels(stencil_blocks) stencil_func_ir.blocks = stencil_blocks return stencil_func_ir, sf.get_return_type(args)[0], arg_to_arr_dict