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 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