def test_inline_update_target_def(self):
def test_impl(a):
if a == 1:
b = 2
else:
b = 3
return b
func_ir = compiler.run_frontend(test_impl)
blocks = list(func_ir.blocks.values())
for block in blocks:
for i, stmt in enumerate(block.body):
# match b = 2 and replace with lambda: 2
if (isinstance(stmt, ir.Assign) and isinstance(stmt.value, ir.Var)
and guard(find_const, func_ir, stmt.value) == 2):
# replace expr with a dummy call
func_ir._definitions[stmt.target.name].remove(stmt.value)
stmt.value = ir.Expr.call(ir.Var(block.scope, "myvar", loc=stmt.loc), (), (), stmt.loc)
func_ir._definitions[stmt.target.name].append(stmt.value)
#func = g.py_func#
inline_closure_call(func_ir, {}, block, i, lambda: 2)
break
self.assertEqual(len(func_ir._definitions['b']), 2)
def _run_inliner(
self,
state,
inline_type,
sig,
template,
arg_typs,
expr,
i,
impl,
block,
work_list,
is_method,
):
from numba.core.inline_closurecall import (
inline_closure_call,
callee_ir_validator,
)
do_inline = True
if not inline_type.is_always_inline:
from numba.core.typing.templates import _inline_info
caller_inline_info = _inline_info(
state.func_ir,
state.type_annotation.typemap,
state.type_annotation.calltypes,
sig,
)
# must be a cost-model function, run the function
iinfo = template._inline_overloads[arg_typs]["iinfo"]
if inline_type.has_cost_model:
do_inline = inline_type.value(expr, caller_inline_info, iinfo)
else:
assert "unreachable"
if do_inline:
if is_method:
if not self._add_method_self_arg(state, expr):
return False
arg_typs = template._inline_overloads[arg_typs]["folded_args"]
# pass is typed so use the callee globals
inline_closure_call(
state.func_ir,
impl.__globals__,
block,
i,
impl,
typingctx=state.typingctx,
arg_typs=arg_typs,
typemap=state.type_annotation.typemap,
calltypes=state.type_annotation.calltypes,
work_list=work_list,
replace_freevars=False,
callee_validator=callee_ir_validator,
)
return True
else:
return False
def run_pass(self, state):
# assuming the function has one block with one call inside
assert len(state.func_ir.blocks) == 1
block = list(state.func_ir.blocks.values())[0]
for i, stmt in enumerate(block.body):
if (guard(find_callname, state.func_ir, stmt.value)
is not None):
inline_closure_call(state.func_ir, {}, block, i,
foo.py_func, state.typingctx,
(state.type_annotation.typemap[stmt.value.args[0].name],),
state.type_annotation.typemap, state.calltypes)
break
return True
def run_pass(self, state):
# assuming the function has one block with one call inside
assert len(state.func_ir.blocks) == 1
block = list(state.func_ir.blocks.values())[0]
for i, stmt in enumerate(block.body):
if guard(find_callname,state.func_ir, stmt.value) is not None:
inline_closure_call(state.func_ir, {}, block, i, lambda: None,
state.typingctx, (), state.type_annotation.typemap,
state.type_annotation.calltypes)
break
# also fix up the IR
post_proc = postproc.PostProcessor(state.func_ir)
post_proc.run()
post_proc.remove_dels()
return True
def test_inline_var_dict_ret(self):
# make sure inline_closure_call returns the variable replacement dict
# and it contains the original variable name used in locals
@njit(locals={'b': types.float64})
def g(a):
b = a + 1
return b
def test_impl():
return g(1)
func_ir = compiler.run_frontend(test_impl)
blocks = list(func_ir.blocks.values())
for block in blocks:
for i, stmt in enumerate(block.body):
if (isinstance(stmt, ir.Assign)
and isinstance(stmt.value, ir.Expr)
and stmt.value.op == 'call'):
func_def = guard(get_definition, func_ir, stmt.value.func)
if (isinstance(func_def, (ir.Global, ir.FreeVar))
and isinstance(func_def.value, CPUDispatcher)):
py_func = func_def.value.py_func
_, var_map = inline_closure_call(
func_ir, py_func.__globals__, block, i, py_func)
break
self.assertTrue('b' in var_map)
def _run_inliner(
func_ir,
sig,
template,
arg_typs,
expr,
i,
py_func,
block,
work_list,
typemap,
calltypes,
typingctx,
targetctx,
):
from numba.core.inline_closurecall import (
callee_ir_validator,
inline_closure_call,
)
# pass is typed so use the callee globals
inline_closure_call(
func_ir,
py_func.__globals__,
block,
i,
py_func,
typingctx=typingctx,
targetctx=targetctx,
arg_typs=arg_typs,
typemap=typemap,
calltypes=calltypes,
work_list=work_list,
replace_freevars=False,
callee_validator=callee_ir_validator,
)
return True
def inline_calls(func_ir, _locals):
work_list = list(func_ir.blocks.items())
while work_list:
label, block = work_list.pop()
for i, instr in enumerate(block.body):
if isinstance(instr, ir.Assign):
lhs = instr.target
expr = instr.value
if isinstance(expr, ir.Expr) and expr.op == 'call':
func_def = guard(get_definition, func_ir, expr.func)
if (isinstance(func_def, (ir.Global, ir.FreeVar))
and isinstance(func_def.value, CPUDispatcher)):
py_func = func_def.value.py_func
inline_out = inline_closure_call(func_ir,
py_func.__globals__,
block,
i,
py_func,
work_list=work_list)
# TODO remove if when inline_closure_call() output fix
# is merged in Numba
if isinstance(inline_out, tuple):
var_dict = inline_out[1]
# TODO: update '##distributed' and '##threaded' in _locals
_locals.update(
(var_dict[k].name, v)
for k, v in func_def.value.locals.items()
if k in var_dict)
# for block in new_blocks:
# work_list.append(block)
# current block is modified, skip the rest
# (included in new blocks)
break
# sometimes type inference fails after inlining since blocks are inserted
# at the end and there are agg constraints (categorical_split case)
# CFG simplification fixes this case
func_ir.blocks = ir_utils.simplify_CFG(func_ir.blocks)
