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