def _loop_lift_modify_blocks(bytecode, loopinfo, blocks,
typingctx, targetctx, flags, locals):
"""
Modify the block inplace to call to the lifted-loop.
Returns a dictionary of blocks of the lifted-loop.
"""
from numba.dispatcher import LiftedLoop
# Copy loop blocks
loop = loopinfo.loop
loopblockkeys = set(loop.body) | set(loop.entries) | set(loop.exits)
loopblocks = dict((k, blocks[k].copy()) for k in loopblockkeys)
# Modify the loop blocks
_loop_lift_prepare_loop_func(loopinfo, loopblocks)
# Create an intrepreter for the lifted loop
interp = Interpreter.from_blocks(bytecode=bytecode, blocks=loopblocks,
override_args=loopinfo.inputs,
force_non_generator=True)
liftedloop = LiftedLoop(interp, typingctx, targetctx, flags, locals)
# modify for calling into liftedloop
callblock = _loop_lift_modify_call_block(liftedloop, blocks[loopinfo.callfrom],
loopinfo.inputs, loopinfo.outputs,
loopinfo.returnto)
# remove blocks
for k in loopblockkeys:
del blocks[k]
# update main interpreter callsite into the liftedloop
blocks[loopinfo.callfrom] = callblock
return liftedloop
def loop_lifting(interp, typingctx, targetctx, flags, locals):
"""
Loop lifting transformation.
Given a interpreter `interp` returns a 2 tuple of
`(toplevel_interp, [loop0_interp, loop1_interp, ....])`
"""
blocks = interp.blocks.copy()
cfg = compute_cfg_from_blocks(blocks)
loopinfos = _loop_lift_get_candidate_infos(cfg, blocks,
interp.variable_lifetime.livemap)
loops = []
for loopinfo in loopinfos:
lifted = _loop_lift_modify_blocks(interp.bytecode, loopinfo, blocks,
typingctx, targetctx, flags, locals)
loops.append(lifted)
# make main interpreter
main = Interpreter.from_blocks(bytecode=interp.bytecode,
blocks=blocks,
used_globals=interp.used_globals)
return main, loops
def get_func_ir(func):
func_id = FunctionIdentity.from_function(func)
bc = ByteCode(func_id=func_id)
interp = Interpreter(func_id)
func_ir = interp.interpret(bc)
return func_ir