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 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 stage_inline_test_pass(self):
# assuming the function has one block with one call inside
assert len(self.func_ir.blocks) == 1
block = list(self.func_ir.blocks.values())[0]
for i, stmt in enumerate(block.body):
if guard(find_callname,self.func_ir, stmt.value) is not None:
inline_closure_call(self.func_ir, {}, block, i, lambda: None,
self.typingctx, (), self.typemap, self.calltypes)
break
def stage_inline_test_pass(self):
# assuming the function has one block with one call inside
assert len(self.func_ir.blocks) == 1
block = list(self.func_ir.blocks.values())[0]
for i, stmt in enumerate(block.body):
if guard(find_callname, self.func_ir, stmt.value) is not None:
inline_closure_call(self.func_ir, {}, block, i, lambda: None,
self.typingctx, (), self.typemap,
self.calltypes)
break
def _run_inliner(
self,
state,
inline_type,
sig,
template,
arg_typs,
expr,
i,
impl,
block,
work_list,
is_method,
):
from numba.inline_closurecall import (inline_closure_call,
callee_ir_validator)
do_inline = True
if not inline_type.is_always_inline:
from numba.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 stage_inline_test_pass(self):
# assuming the function has one block with one call inside
assert len(self.func_ir.blocks) == 1
block = list(self.func_ir.blocks.values())[0]
for i, stmt in enumerate(block.body):
if guard(find_callname,self.func_ir, stmt.value) is not None:
inline_closure_call(self.func_ir, {}, block, i, lambda: None,
self.typingctx, (), self.typemap, self.calltypes)
# also fix up the IR so that ir.Dels appear correctly/in correct
# locations
post_proc = numba.postproc.PostProcessor(self.func_ir)
post_proc.run()
break
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 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
@numba.njit(locals={'b': numba.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 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)
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
@numba.njit(locals={'b': numba.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 inline_calls(func_ir):
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) and isinstance(
func_def.value, CPUDispatcher):
py_func = func_def.value.py_func
new_blocks = inline_closure_call(
func_ir,
func_ir.func_id.func.__globals__,
block,
i,
py_func,
work_list=work_list)
# for block in new_blocks:
# work_list.append(block)
# current block is modified, skip the rest
# (included in new blocks)
break
def _do_work(self, state, work_list, block, i, expr):
from numba.inline_closurecall import (inline_closure_call,
callee_ir_validator)
from numba.compiler import run_frontend
from numba.targets.cpu import InlineOptions
# try and get a definition for the call, this isn't always possible as
# it might be a eval(str)/part generated awaiting update etc. (parfors)
to_inline = None
try:
to_inline = state.func_ir.get_definition(expr.func)
except Exception:
if self._DEBUG:
print("Cannot find definition for %s" % expr.func)
return False
# do not handle closure inlining here, another pass deals with that.
if getattr(to_inline, 'op', False) == 'make_function':
return False
# see if the definition is a "getattr", in which case walk the IR to
# try and find the python function via the module from which it's
# imported, this should all be encoded in the IR.
if getattr(to_inline, 'op', False) == 'getattr':
val = resolve_func_from_module(state.func_ir, to_inline)
else:
# This is likely a freevar or global
#
# NOTE: getattr 'value' on a call may fail if it's an ir.Expr as
# getattr is overloaded to look in _kws.
try:
val = getattr(to_inline, 'value', False)
except Exception:
raise GuardException
# if something was found...
if val:
# check it's dispatcher-like, the targetoptions attr holds the
# kwargs supplied in the jit decorator and is where 'inline' will
# be if it is present.
topt = getattr(val, 'targetoptions', False)
if topt:
inline_type = topt.get('inline', None)
# has 'inline' been specified?
if inline_type is not None:
inline_opt = InlineOptions(inline_type)
# Could this be inlinable?
if not inline_opt.is_never_inline:
# yes, it could be inlinable
do_inline = True
pyfunc = val.py_func
# Has it got an associated cost model?
if inline_opt.has_cost_model:
# yes, it has a cost model, use it to determine
# whether to do the inline
py_func_ir = run_frontend(pyfunc)
do_inline = inline_type(state.func_ir, py_func_ir)
# if do_inline is True then inline!
if do_inline:
inline_closure_call(
state.func_ir,
pyfunc.__globals__,
block,
i,
pyfunc,
work_list=work_list,
callee_validator=callee_ir_validator)
return True
return False
def _do_work(self, state, work_list, block, i, expr):
from numba.inline_closurecall import (inline_closure_call,
callee_ir_validator)
# try and get a definition for the call, this isn't always possible as
# it might be a eval(str)/part generated awaiting update etc. (parfors)
to_inline = None
try:
to_inline = state.func_ir.get_definition(expr.func)
except Exception:
return False
# do not handle closure inlining here, another pass deals with that.
if getattr(to_inline, 'op', False) == 'make_function':
return False
# check this is a known and typed function
try:
func_ty = state.type_annotation.typemap[expr.func.name]
except KeyError:
# e.g. Calls to CUDA Intrinsic have no mapped type so KeyError
return False
if not hasattr(func_ty, 'get_call_type'):
return False
# search the templates for this overload looking for "inline"
templates = getattr(func_ty, 'templates', None)
if templates is None:
return False
sig = state.type_annotation.calltypes[expr]
impl = None
for template in templates:
inline_type = getattr(template, '_inline', None)
if inline_type is None:
# inline not defined
continue
if not inline_type.is_never_inline:
try:
impl = template._overload_func(*sig.args)
if impl is None:
raise Exception # abort for this template
break
except Exception:
continue
else:
return False
# at this point we know we maybe want to inline something and there's
# definitely something that could be inlined.
do_inline = True
if not inline_type.is_always_inline:
from numba.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[sig.args]['iinfo']
if inline_type.has_cost_model:
do_inline = inline_type.value(expr, caller_inline_info, iinfo)
else:
assert 'unreachable'
if do_inline:
arg_typs = template._inline_overloads[sig.args]['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
