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)