from collections import defaultdict, namedtuple from contextlib import contextmanager from copy import deepcopy, copy import warnings from numba.core.compiler_machinery import (FunctionPass, AnalysisPass, SSACompliantMixin, register_pass) from numba.core import (errors, types, ir, bytecode, postproc, rewrites, config, transforms) from numba.misc.special import literal_unroll from numba.core.analysis import (dead_branch_prune, rewrite_semantic_constants, find_literally_calls, compute_cfg_from_blocks, compute_use_defs) from numba.core.ir_utils import (guard, resolve_func_from_module, simplify_CFG, GuardException, convert_code_obj_to_function, mk_unique_var, build_definitions, replace_var_names, get_name_var_table, compile_to_numba_ir, get_definition, find_max_label, rename_labels, transfer_scope, fixup_var_define_in_scope, ) from numba.core.ssa import reconstruct_ssa from numba.core import interpreter @contextmanager def fallback_context(state, msg): """ Wraps code that would signal a fallback to object mode """ try: yield except Exception as e: if not state.status.can_fallback: raise else: # Clear all references attached to the traceback e = e.with_traceback(None) # this emits a warning containing the error message body in the # case of fallback from npm to objmode loop_lift = '' if state.flags.enable_looplift else 'OUT' msg_rewrite = ("\nCompilation is falling back to object mode " "WITH%s looplifting enabled because %s" % (loop_lift, msg)) warnings.warn_explicit('%s due to: %s' % (msg_rewrite, e), errors.NumbaWarning, state.func_id.filename, state.func_id.firstlineno) raise @register_pass(mutates_CFG=True, analysis_only=False) class ExtractByteCode(FunctionPass): _name = "extract_bytecode" def __init__(self): FunctionPass.__init__(self) def run_pass(self, state): """ Extract bytecode from function """ func_id = state['func_id'] bc = bytecode.ByteCode(func_id) if config.DUMP_BYTECODE: print(bc.dump()) state['bc'] = bc return True @register_pass(mutates_CFG=True, analysis_only=False) class TranslateByteCode(FunctionPass): _name = "translate_bytecode" def __init__(self): FunctionPass.__init__(self) def run_pass(self, state): """ Analyze bytecode and translating to Numba IR """ func_id = state['func_id'] bc = state['bc'] interp = interpreter.Interpreter(func_id) func_ir = interp.interpret(bc) state["func_ir"] = func_ir return True @register_pass(mutates_CFG=True, analysis_only=False) class FixupArgs(FunctionPass): _name = "fixup_args" def __init__(self): FunctionPass.__init__(self) def run_pass(self, state): state['nargs'] = state['func_ir'].arg_count if not state['args'] and state['flags'].force_pyobject: # Allow an empty argument types specification when object mode # is explicitly requested. state['args'] = (types.pyobject,) * state['nargs'] elif len(state['args']) != state['nargs']: raise TypeError("Signature mismatch: %d argument types given, " "but function takes %d arguments" % (len(state['args']), state['nargs'])) return True @register_pass(mutates_CFG=True, analysis_only=False) class IRProcessing(FunctionPass): _name = "ir_processing" def __init__(self): FunctionPass.__init__(self) def run_pass(self, state): func_ir = state['func_ir'] post_proc = postproc.PostProcessor(func_ir) post_proc.run() if config.DEBUG or config.DUMP_IR: name = func_ir.func_id.func_qualname print(("IR DUMP: %s" % name).center(80, "-")) func_ir.dump() if func_ir.is_generator: print(("GENERATOR INFO: %s" % name).center(80, "-")) func_ir.dump_generator_info() return True @register_pass(mutates_CFG=True, analysis_only=False) class RewriteSemanticConstants(FunctionPass): _name = "rewrite_semantic_constants" def __init__(self): FunctionPass.__init__(self) def run_pass(self, state): """ This prunes dead branches, a dead branch is one which is derivable as not taken at compile time purely based on const/literal evaluation. """ assert state.func_ir msg = ('Internal error in pre-inference dead branch pruning ' 'pass encountered during compilation of ' 'function "%s"' % (state.func_id.func_name,)) with fallback_context(state, msg): rewrite_semantic_constants(state.func_ir, state.args) return True @register_pass(mutates_CFG=True, analysis_only=False) class DeadBranchPrune(SSACompliantMixin, FunctionPass): _name = "dead_branch_prune" def __init__(self): FunctionPass.__init__(self) def run_pass(self, state): """ This prunes dead branches, a dead branch is one which is derivable as not taken at compile time purely based on const/literal evaluation. """ # purely for demonstration purposes, obtain the analysis from a pass # declare as a required dependent semantic_const_analysis = self.get_analysis(type(self)) # noqa assert state.func_ir msg = ('Internal error in pre-inference dead branch pruning ' 'pass encountered during compilation of ' 'function "%s"' % (state.func_id.func_name,)) with fallback_context(state, msg): dead_branch_prune(state.func_ir, state.args) return True def get_analysis_usage(self, AU): AU.add_required(RewriteSemanticConstants) @register_pass(mutates_CFG=True, analysis_only=False) class InlineClosureLikes(FunctionPass): _name = "inline_closure_likes" def __init__(self): FunctionPass.__init__(self) def run_pass(self, state): # Ensure we have an IR and type information. assert state.func_ir # if the return type is a pyobject, there's no type info available and # no ability to resolve certain typed function calls in the array # inlining code, use this variable to indicate typed_pass = not isinstance(state.return_type, types.misc.PyObject) from numba.core.inline_closurecall import InlineClosureCallPass inline_pass = InlineClosureCallPass( state.func_ir, state.flags.auto_parallel, state.parfor_diagnostics.replaced_fns, typed_pass) inline_pass.run() # Remove all Dels, and re-run postproc post_proc = postproc.PostProcessor(state.func_ir) post_proc.run() fixup_var_define_in_scope(state.func_ir.blocks) return True @register_pass(mutates_CFG=True, analysis_only=False) class GenericRewrites(FunctionPass): _name = "generic_rewrites" def __init__(self): FunctionPass.__init__(self) def run_pass(self, state): """ Perform any intermediate representation rewrites before type inference. """ assert state.func_ir msg = ('Internal error in pre-inference rewriting ' 'pass encountered during compilation of ' 'function "%s"' % (state.func_id.func_name,)) with fallback_context(state, msg): rewrites.rewrite_registry.apply('before-inference', state) return True @register_pass(mutates_CFG=True, analysis_only=False) class WithLifting(FunctionPass): _name = "with_lifting" def __init__(self): FunctionPass.__init__(self) def run_pass(self, state): """ Extract with-contexts """ main, withs = transforms.with_lifting( func_ir=state.func_ir, typingctx=state.typingctx, targetctx=state.targetctx, flags=state.flags, locals=state.locals, ) if withs: from numba.core.compiler import compile_ir, _EarlyPipelineCompletion cres = compile_ir(state.typingctx, state.targetctx, main, state.args, state.return_type, state.flags, state.locals, lifted=tuple(withs), lifted_from=None, pipeline_class=type(state.pipeline)) raise _EarlyPipelineCompletion(cres) return True @register_pass(mutates_CFG=True, analysis_only=False) class InlineInlinables(FunctionPass): """ This pass will inline a function wrapped by the numba.jit decorator directly into the site of its call depending on the value set in the 'inline' kwarg to the decorator. This is an untyped pass. CFG simplification is performed at the end of the pass but no block level clean up is performed on the mutated IR (typing information is not available to do so). """ _name = "inline_inlinables" _DEBUG = False def __init__(self): FunctionPass.__init__(self) def run_pass(self, state): """Run inlining of inlinables """ if self._DEBUG: print('before inline'.center(80, '-')) print(state.func_ir.dump()) print(''.center(80, '-')) from numba.core.inline_closurecall import (InlineWorker, callee_ir_validator) inline_worker = InlineWorker(state.typingctx, state.targetctx, state.locals, state.pipeline, state.flags, validator=callee_ir_validator) modified = False # use a work list, look for call sites via `ir.Expr.op == call` and # then pass these to `self._do_work` to make decisions about inlining. work_list = list(state.func_ir.blocks.items()) while work_list: label, block = work_list.pop() for i, instr in enumerate(block.body): if isinstance(instr, ir.Assign): expr = instr.value if isinstance(expr, ir.Expr) and expr.op == 'call': if guard(self._do_work, state, work_list, block, i, expr, inline_worker): modified = True break # because block structure changed if modified: # clean up unconditional branches that appear due to inlined # functions introducing blocks cfg = compute_cfg_from_blocks(state.func_ir.blocks) for dead in cfg.dead_nodes(): del state.func_ir.blocks[dead] post_proc = postproc.PostProcessor(state.func_ir) post_proc.run() state.func_ir.blocks = simplify_CFG(state.func_ir.blocks) if self._DEBUG: print('after inline'.center(80, '-')) print(state.func_ir.dump()) print(''.center(80, '-')) return True def _do_work(self, state, work_list, block, i, expr, inline_worker): from numba.core.compiler import run_frontend from numba.core.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(expr, state.func_ir, py_func_ir) # if do_inline is True then inline! if do_inline: _, _, _, new_blocks = \ inline_worker.inline_function(state.func_ir, block, i, pyfunc,) if work_list is not None: for blk in new_blocks: work_list.append(blk) return True return False @register_pass(mutates_CFG=False, analysis_only=False) class PreserveIR(AnalysisPass): """ Preserves the IR in the metadata """ _name = "preserve_ir" def __init__(self): AnalysisPass.__init__(self) def run_pass(self, state): state.metadata['preserved_ir'] = state.func_ir.copy() return False @register_pass(mutates_CFG=False, analysis_only=True) class FindLiterallyCalls(FunctionPass): """Find calls to `numba.literally()` and signal if its requirement is not satisfied. """ _name = "find_literally" def __init__(self): FunctionPass.__init__(self) def run_pass(self, state): find_literally_calls(state.func_ir, state.args) return False @register_pass(mutates_CFG=True, analysis_only=False) class CanonicalizeLoopExit(FunctionPass): """A pass to canonicalize loop exit by splitting it from function exit. """ _name = "canonicalize_loop_exit" def __init__(self): FunctionPass.__init__(self) def run_pass(self, state): fir = state.func_ir cfg = compute_cfg_from_blocks(fir.blocks) status = False for loop in cfg.loops().values(): for exit_label in loop.exits: if exit_label in cfg.exit_points(): self._split_exit_block(fir, cfg, exit_label) status = True fir._reset_analysis_variables() vlt = postproc.VariableLifetime(fir.blocks) fir.variable_lifetime = vlt return status def _split_exit_block(self, fir, cfg, exit_label): curblock = fir.blocks[exit_label] newlabel = exit_label + 1 newlabel = find_max_label(fir.blocks) + 1 fir.blocks[newlabel] = curblock newblock = ir.Block(scope=curblock.scope, loc=curblock.loc) newblock.append(ir.Jump(newlabel, loc=curblock.loc)) fir.blocks[exit_label] = newblock # Rename all labels fir.blocks = rename_labels(fir.blocks) @register_pass(mutates_CFG=True, analysis_only=False) class CanonicalizeLoopEntry(FunctionPass): """A pass to canonicalize loop header by splitting it from function entry. This is needed for loop-lifting; esp in py3.8 """ _name = "canonicalize_loop_entry" def __init__(self): FunctionPass.__init__(self) def run_pass(self, state): fir = state.func_ir cfg = compute_cfg_from_blocks(fir.blocks) status = False for loop in cfg.loops().values(): if len(loop.entries) == 1: [entry_label] = loop.entries if entry_label == cfg.entry_point(): self._split_entry_block(fir, cfg, loop, entry_label) status = True fir._reset_analysis_variables() vlt = postproc.VariableLifetime(fir.blocks) fir.variable_lifetime = vlt return status def _split_entry_block(self, fir, cfg, loop, entry_label): # Find iterator inputs into the for-loop header header_block = fir.blocks[loop.header] deps = set() for expr in header_block.find_exprs(op="iternext"): deps.add(expr.value) # Find the getiter for each iterator entry_block = fir.blocks[entry_label] # Find the start of loop entry statement that needs to be included. startpt = None list_of_insts = list(entry_block.find_insts(ir.Assign)) for assign in reversed(list_of_insts): if assign.target in deps: rhs = assign.value if isinstance(rhs, ir.Var): if rhs.is_temp: deps.add(rhs) elif isinstance(rhs, ir.Expr): expr = rhs if expr.op == 'getiter': startpt = assign if expr.value.is_temp: deps.add(expr.value) elif expr.op == 'call': defn = guard(get_definition, fir, expr.func) if isinstance(defn, ir.Global): if expr.func.is_temp: deps.add(expr.func) elif isinstance(rhs, ir.Global) and rhs.value is range: startpt = assign if startpt is None: return splitpt = entry_block.body.index(startpt) new_block = entry_block.copy() new_block.body = new_block.body[splitpt:] new_block.loc = new_block.body[0].loc new_label = find_max_label(fir.blocks) + 1 entry_block.body = entry_block.body[:splitpt] entry_block.append(ir.Jump(new_label, loc=new_block.loc)) fir.blocks[new_label] = new_block # Rename all labels fir.blocks = rename_labels(fir.blocks) @register_pass(mutates_CFG=False, analysis_only=True) class PrintIRCFG(FunctionPass): _name = "print_ir_cfg" def __init__(self): FunctionPass.__init__(self) self._ver = 0 def run_pass(self, state): fir = state.func_ir self._ver += 1 fir.render_dot(filename_prefix='v{}'.format(self._ver)).render() return False @register_pass(mutates_CFG=True, analysis_only=False) class MakeFunctionToJitFunction(FunctionPass): """ This swaps an ir.Expr.op == "make_function" i.e. a closure, for a compiled function containing the closure body and puts it in ir.Global. It's a 1:1 statement value swap. `make_function` is already untyped """ _name = "make_function_op_code_to_jit_function" def __init__(self): FunctionPass.__init__(self) def run_pass(self, state): from numba import njit func_ir = state.func_ir mutated = False for idx, blk in func_ir.blocks.items(): for stmt in blk.body: if isinstance(stmt, ir.Assign): if isinstance(stmt.value, ir.Expr): if stmt.value.op == "make_function": node = stmt.value getdef = func_ir.get_definition kw_default = getdef(node.defaults) ok = False if (kw_default is None or isinstance(kw_default, ir.Const)): ok = True elif isinstance(kw_default, tuple): ok = all([isinstance(getdef(x), ir.Const) for x in kw_default]) elif isinstance(kw_default, ir.Expr): if kw_default.op != "build_tuple": continue ok = all([isinstance(getdef(x), ir.Const) for x in kw_default.items]) if not ok: continue pyfunc = convert_code_obj_to_function(node, func_ir) func = njit()(pyfunc) new_node = ir.Global(node.code.co_name, func, stmt.loc) stmt.value = new_node mutated |= True # if a change was made the del ordering is probably wrong, patch up if mutated: post_proc = postproc.PostProcessor(func_ir) post_proc.run() return mutated @register_pass(mutates_CFG=True, analysis_only=False) class TransformLiteralUnrollConstListToTuple(FunctionPass): """ This pass spots a `literal_unroll([])` and rewrites it as a `literal_unroll(tuple())`. """ _name = "transform_literal_unroll_const_list_to_tuple" _accepted_types = (types.BaseTuple, types.LiteralList) def __init__(self): FunctionPass.__init__(self) def run_pass(self, state): mutated = False func_ir = state.func_ir for label, blk in func_ir.blocks.items(): calls = [_ for _ in blk.find_exprs('call')] for call in calls: glbl = guard(get_definition, func_ir, call.func) if glbl and isinstance(glbl, (ir.Global, ir.FreeVar)): # find a literal_unroll if glbl.value is literal_unroll: if len(call.args) > 1: msg = "literal_unroll takes one argument, found %s" raise errors.UnsupportedError(msg % len(call.args), call.loc) # get the arg, make sure its a build_list unroll_var = call.args[0] to_unroll = guard(get_definition, func_ir, unroll_var) if (isinstance(to_unroll, ir.Expr) and to_unroll.op == "build_list"): # make sure they are all const items in the list for i, item in enumerate(to_unroll.items): val = guard(get_definition, func_ir, item) if not val: msg = ("multiple definitions for variable " "%s, cannot resolve constant") raise errors.UnsupportedError(msg % item, to_unroll.loc) if not isinstance(val, ir.Const): msg = ("Found non-constant value at " "position %s in a list argument to " "literal_unroll" % i) raise errors.UnsupportedError(msg, to_unroll.loc) # The above appears ok, now swap the build_list for # a built tuple. # find the assignment for the unroll target to_unroll_lhs = guard(get_definition, func_ir, unroll_var, lhs_only=True) if to_unroll_lhs is None: msg = ("multiple definitions for variable " "%s, cannot resolve constant") raise errors.UnsupportedError(msg % unroll_var, to_unroll.loc) # scan all blocks looking for the LHS for b in func_ir.blocks.values(): asgn = b.find_variable_assignment( to_unroll_lhs.name) if asgn is not None: break else: msg = ("Cannot find assignment for known " "variable %s") % to_unroll_lhs.name raise errors.CompilerError(msg, to_unroll.loc) # Create a tuple with the list items as contents tup = ir.Expr.build_tuple(to_unroll.items, to_unroll.loc) # swap the list for the tuple asgn.value = tup mutated = True elif (isinstance(to_unroll, ir.Expr) and to_unroll.op == "build_tuple"): # this is fine, do nothing pass elif (isinstance(to_unroll, (ir.Global, ir.FreeVar)) and isinstance(to_unroll.value, tuple)): # this is fine, do nothing pass elif isinstance(to_unroll, ir.Arg): # this is only fine if the arg is a tuple ty = state.typemap[to_unroll.name] if not isinstance(ty, self._accepted_types): msg = ("Invalid use of literal_unroll with a " "function argument, only tuples are " "supported as function arguments, found " "%s") % ty raise errors.UnsupportedError(msg, to_unroll.loc) else: extra = None if isinstance(to_unroll, ir.Expr): # probably a slice if to_unroll.op == "getitem": ty = state.typemap[to_unroll.value.name] # check if this is a tuple slice if not isinstance(ty, self._accepted_types): extra = "operation %s" % to_unroll.op loc = to_unroll.loc elif isinstance(to_unroll, ir.Arg): extra = "non-const argument %s" % to_unroll.name loc = to_unroll.loc else: if to_unroll is None: extra = ('multiple definitions of ' 'variable "%s".' % unroll_var.name) loc = unroll_var.loc else: loc = to_unroll.loc extra = "unknown problem" if extra: msg = ("Invalid use of literal_unroll, " "argument should be a tuple or a list " "of constant values. Failure reason: " "found %s" % extra) raise errors.UnsupportedError(msg, loc) return mutated @register_pass(mutates_CFG=True, analysis_only=False) class MixedContainerUnroller(FunctionPass): _name = "mixed_container_unroller" _DEBUG = False _accepted_types = (types.BaseTuple, types.LiteralList) def __init__(self): FunctionPass.__init__(self) def analyse_tuple(self, tup): """ Returns a map of type->list(indexes) for a typed tuple """ d = defaultdict(list) for i, ty in enumerate(tup): d[ty].append(i) return d def add_offset_to_labels_w_ignore(self, blocks, offset, ignore=None): """add an offset to all block labels and jump/branch targets don't add an offset to anything in the ignore list """ if ignore is None: ignore = set() new_blocks = {} for l, b in blocks.items(): # some parfor last blocks might be empty term = None if b.body: term = b.body[-1] if isinstance(term, ir.Jump): if term.target not in ignore: b.body[-1] = ir.Jump(term.target + offset, term.loc) if isinstance(term, ir.Branch): if term.truebr not in ignore: new_true = term.truebr + offset else: new_true = term.truebr if term.falsebr not in ignore: new_false = term.falsebr + offset else: new_false = term.falsebr b.body[-1] = ir.Branch(term.cond, new_true, new_false, term.loc) new_blocks[l + offset] = b return new_blocks def inject_loop_body(self, switch_ir, loop_ir, caller_max_label, dont_replace, switch_data): """ Injects the "loop body" held in `loop_ir` into `switch_ir` where ever there is a statement of the form `SENTINEL. = RHS`. It also: * Finds and then deliberately does not relabel non-local jumps so as to make the switch table suitable for injection into the IR from which the loop body was derived. * Looks for `typed_getitem` and wires them up to loop body version specific variables or, if possible, directly writes in their constant value at their use site. Args: - switch_ir, the switch table with SENTINELS as generated by self.gen_switch - loop_ir, the IR of the loop blocks (derived from the original func_ir) - caller_max_label, the maximum label in the func_ir caller - dont_replace, variables that should not be renamed (to handle references to variables that are incoming at the loop head/escaping at the loop exit. - switch_data, the switch table data used to generated the switch_ir, can be generated by self.analyse_tuple. Returns: - A type specific switch table with each case containing a versioned loop body suitable for injection as a replacement for the loop_ir. """ # Switch IR came from code gen, immediately relabel to prevent # collisions with IR derived from the user code (caller) switch_ir.blocks = self.add_offset_to_labels_w_ignore( switch_ir.blocks, caller_max_label + 1) # Find the sentinels and validate the form sentinel_exits = set() sentinel_blocks = [] for lbl, blk in switch_ir.blocks.items(): for i, stmt in enumerate(blk.body): if isinstance(stmt, ir.Assign): if "SENTINEL" in stmt.target.name: sentinel_blocks.append(lbl) sentinel_exits.add(blk.body[-1].target) break assert len(sentinel_exits) == 1 # should only be 1 exit switch_ir.blocks.pop(sentinel_exits.pop()) # kill the exit, it's dead # find jumps that are non-local, we won't relabel these ignore_set = set() local_lbl = [x for x in loop_ir.blocks.keys()] for lbl, blk in loop_ir.blocks.items(): for i, stmt in enumerate(blk.body): if isinstance(stmt, ir.Jump): if stmt.target not in local_lbl: ignore_set.add(stmt.target) if isinstance(stmt, ir.Branch): if stmt.truebr not in local_lbl: ignore_set.add(stmt.truebr) if stmt.falsebr not in local_lbl: ignore_set.add(stmt.falsebr) # make sure the generated switch table matches the switch data assert len(sentinel_blocks) == len(switch_data) # replace the sentinel_blocks with the loop body for lbl, branch_ty in zip(sentinel_blocks, switch_data.keys()): loop_blocks = deepcopy(loop_ir.blocks) # relabel blocks WRT switch table, each block replacement will shift # the maximum label max_label = max(switch_ir.blocks.keys()) loop_blocks = self.add_offset_to_labels_w_ignore( loop_blocks, max_label + 1, ignore_set) # start label loop_start_lbl = min(loop_blocks.keys()) # fix the typed_getitem locations in the loop blocks for blk in loop_blocks.values(): new_body = [] for stmt in blk.body: if isinstance(stmt, ir.Assign): if (isinstance(stmt.value, ir.Expr) and stmt.value.op == "typed_getitem"): if isinstance(branch_ty, types.Literal): new_const_name = mk_unique_var("branch_const") new_const_var = ir.Var( blk.scope, new_const_name, stmt.loc) new_const_val = ir.Const( branch_ty.literal_value, stmt.loc) const_assign = ir.Assign( new_const_val, new_const_var, stmt.loc) new_assign = ir.Assign( new_const_var, stmt.target, stmt.loc) new_body.append(const_assign) new_body.append(new_assign) dont_replace.append(new_const_name) else: orig = stmt.value new_typed_getitem = ir.Expr.typed_getitem( value=orig.value, dtype=branch_ty, index=orig.index, loc=orig.loc) new_assign = ir.Assign( new_typed_getitem, stmt.target, stmt.loc) new_body.append(new_assign) else: new_body.append(stmt) else: new_body.append(stmt) blk.body = new_body # rename var_table = get_name_var_table(loop_blocks) drop_keys = [] for k, v in var_table.items(): if v.name in dont_replace: drop_keys.append(k) for k in drop_keys: var_table.pop(k) new_var_dict = {} for name, var in var_table.items(): new_var_dict[name] = mk_unique_var(name) replace_var_names(loop_blocks, new_var_dict) # clobber the sentinel body and then stuff in the rest switch_ir.blocks[lbl] = deepcopy(loop_blocks[loop_start_lbl]) remaining_keys = [y for y in loop_blocks.keys()] remaining_keys.remove(loop_start_lbl) for k in remaining_keys: switch_ir.blocks[k] = deepcopy(loop_blocks[k]) if self._DEBUG: print("-" * 80 + "EXIT STUFFER") switch_ir.dump() print("-" * 80) return switch_ir def gen_switch(self, data, index): """ Generates a function with a switch table like def foo(): if PLACEHOLDER_INDEX in (): SENTINEL = None elif PLACEHOLDER_INDEX in (): SENTINEL = None ... else: raise RuntimeError The data is a map of (type : indexes) for example: (int64, int64, float64) might give: {int64: [0, 1], float64: [2]} The index is the index variable for the driving range loop over the mixed tuple. """ elif_tplt = "\n\telif PLACEHOLDER_INDEX in (%s,):\n\t\tSENTINEL = None" # Note regarding the insertion of the garbage/defeat variables below: # These values have been designed and inserted to defeat a specific # behaviour of the cpython optimizer. The optimization was introduced # in Python 3.10. # The URL for the BPO is: # https://bugs.python.org/issue44626 # The code for the optimization can be found at: # https://github.com/python/cpython/blob/d41abe8/Python/compile.c#L7533-L7557 # Essentially the CPython optimizer will inline the exit block under # certain circumstances and thus replace the jump with a return if the # exit block is small enough. This is an issue for unroller, as it # looks for a jump, not a return, when it inserts the generated switch # table. # Part of the condition for this optimization to be applied is that the # exit block not exceed a certain (4 at the time of writing) number of # bytecode instructions. We defeat the optimizer by inserting a # sufficient number of instructions so that the exit block is big # enough. We don't care about this garbage, because the generated exit # block is discarded anyway when we smash the switch table into the # original function and so all the inserted garbage is dropped again. # The final lines of the stacktrace w/o this will look like: # # File "/numba/numba/core/untyped_passes.py", line 830, \ # in inject_loop_body # sentinel_exits.add(blk.body[-1].target) # AttributeError: Failed in nopython mode pipeline \ # (step: handles literal_unroll) # Failed in literal_unroll_subpipeline mode pipeline \ # (step: performs mixed container unroll) # 'Return' object has no attribute 'target' # # Which indicates that a Return has been found instead of a Jump b = ('def foo():\n\tif PLACEHOLDER_INDEX in (%s,):\n\t\t' 'SENTINEL = None\n%s\n\telse:\n\t\t' 'raise RuntimeError("Unreachable")\n\t' 'py310_defeat1 = 1\n\t' 'py310_defeat2 = 2\n\t' 'py310_defeat3 = 3\n\t' 'py310_defeat4 = 4\n\t' ) keys = [k for k in data.keys()] elifs = [] for i in range(1, len(keys)): elifs.append(elif_tplt % ','.join(map(str, data[keys[i]]))) src = b % (','.join(map(str, data[keys[0]])), ''.join(elifs)) wstr = src l = {} exec(wstr, {}, l) bfunc = l['foo'] branches = compile_to_numba_ir(bfunc, {}) for lbl, blk in branches.blocks.items(): for stmt in blk.body: if isinstance(stmt, ir.Assign): if isinstance(stmt.value, ir.Global): if stmt.value.name == "PLACEHOLDER_INDEX": stmt.value = index return branches def apply_transform(self, state): # compute new CFG func_ir = state.func_ir cfg = compute_cfg_from_blocks(func_ir.blocks) # find loops loops = cfg.loops() # 0. Find the loops containing literal_unroll and store this # information literal_unroll_info = dict() unroll_info = namedtuple( "unroll_info", [ "loop", "call", "arg", "getitem"]) def get_call_args(init_arg, want): # Chases the assignment of a called value back through a specific # call to a global function "want" and returns the arguments # supplied to that function's call some_call = get_definition(func_ir, init_arg) if not isinstance(some_call, ir.Expr): raise GuardException if not some_call.op == "call": raise GuardException the_global = get_definition(func_ir, some_call.func) if not isinstance(the_global, ir.Global): raise GuardException if the_global.value is not want: raise GuardException return some_call for lbl, loop in loops.items(): # TODO: check the loop head has literal_unroll, if it does but # does not conform to the following then raise # scan loop header iternexts = [_ for _ in func_ir.blocks[loop.header].find_exprs('iternext')] if len(iternexts) != 1: return False for iternext in iternexts: # Walk the canonicalised loop structure and check it # Check loop form range(literal_unroll(container))) phi = guard(get_definition, func_ir, iternext.value) if phi is None: continue # check call global "range" range_call = guard(get_call_args, phi.value, range) if range_call is None: continue range_arg = range_call.args[0] # check call global "len" len_call = guard(get_call_args, range_arg, len) if len_call is None: continue len_arg = len_call.args[0] # check literal_unroll literal_unroll_call = guard(get_definition, func_ir, len_arg) if literal_unroll_call is None: continue if not isinstance(literal_unroll_call, ir.Expr): continue if literal_unroll_call.op != "call": continue literal_func = getattr(literal_unroll_call, 'func', None) if not literal_func: continue call_func = guard(get_definition, func_ir, literal_unroll_call.func) if call_func is None: continue call_func = call_func.value if call_func is literal_unroll: assert len(literal_unroll_call.args) == 1 arg = literal_unroll_call.args[0] typemap = state.typemap resolved_arg = guard(get_definition, func_ir, arg, lhs_only=True) ty = typemap[resolved_arg.name] assert isinstance(ty, self._accepted_types) # loop header is spelled ok, now make sure the body # actually contains a getitem # find a "getitem" tuple_getitem = None for lbl in loop.body: blk = func_ir.blocks[lbl] for stmt in blk.body: if isinstance(stmt, ir.Assign): if (isinstance(stmt.value, ir.Expr) and stmt.value.op == "getitem"): # check for something like a[i] if stmt.value.value != arg: # that failed, so check for the # definition dfn = guard(get_definition, func_ir, stmt.value.value) if dfn is None: continue try: args = getattr(dfn, 'args', False) except KeyError: continue if not args: continue if not args[0] == arg: continue target_ty = state.typemap[arg.name] if not isinstance(target_ty, self._accepted_types): continue tuple_getitem = stmt break if tuple_getitem: break else: continue # no getitem in this loop ui = unroll_info(loop, literal_unroll_call, arg, tuple_getitem) literal_unroll_info[lbl] = ui if not literal_unroll_info: return False # 1. Validate loops, must not have any calls to literal_unroll for test_lbl, test_loop in literal_unroll_info.items(): for ref_lbl, ref_loop in literal_unroll_info.items(): if test_lbl == ref_lbl: # comparing to self! skip continue if test_loop.loop.header in ref_loop.loop.body: msg = ("Nesting of literal_unroll is unsupported") loc = func_ir.blocks[test_loop.loop.header].loc raise errors.UnsupportedError(msg, loc) # 2. Do the unroll, get a loop and process it! lbl, info = literal_unroll_info.popitem() self.unroll_loop(state, info) # 3. Rebuild the state, the IR has taken a hammering func_ir.blocks = simplify_CFG(func_ir.blocks) post_proc = postproc.PostProcessor(func_ir) post_proc.run() if self._DEBUG: print('-' * 80 + "END OF PASS, SIMPLIFY DONE") func_ir.dump() func_ir._definitions = build_definitions(func_ir.blocks) return True def unroll_loop(self, state, loop_info): # The general idea here is to: # 1. Find *a* getitem that conforms to the literal_unroll semantic, # i.e. one that is targeting a tuple with a loop induced index # 2. Compute a structure from the tuple that describes which # iterations of a loop will have which type # 3. Generate a switch table in IR form for the structure in 2 # 4. Switch out getitems for the tuple for a `typed_getitem` # 5. Inject switch table as replacement loop body # 6. Patch up func_ir = state.func_ir getitem_target = loop_info.arg target_ty = state.typemap[getitem_target.name] assert isinstance(target_ty, self._accepted_types) # 1. find a "getitem" that conforms tuple_getitem = [] for lbl in loop_info.loop.body: blk = func_ir.blocks[lbl] for stmt in blk.body: if isinstance(stmt, ir.Assign): if isinstance(stmt.value, ir.Expr) and stmt.value.op == "getitem": # try a couple of spellings... a[i] and ref(a)[i] if stmt.value.value != getitem_target: dfn = func_ir.get_definition(stmt.value.value) try: args = getattr(dfn, 'args', False) except KeyError: continue if not args: continue if not args[0] == getitem_target: continue target_ty = state.typemap[getitem_target.name] if not isinstance(target_ty, self._accepted_types): continue tuple_getitem.append(stmt) if not tuple_getitem: msg = ("Loop unrolling analysis has failed, there's no getitem " "in loop body that conforms to literal_unroll " "requirements.") LOC = func_ir.blocks[loop_info.loop.header].loc raise errors.CompilerError(msg, LOC) # 2. get switch data switch_data = self.analyse_tuple(target_ty) # 3. generate switch IR index = func_ir._definitions[tuple_getitem[0].value.index.name][0] branches = self.gen_switch(switch_data, index) # 4. swap getitems for a typed_getitem, these are actually just # placeholders at this point. When the loop is duplicated they can # be swapped for a typed_getitem of the correct type or if the item # is literal it can be shoved straight into the duplicated loop body for item in tuple_getitem: old = item.value new = ir.Expr.typed_getitem( old.value, types.void, old.index, old.loc) item.value = new # 5. Inject switch table # Find the actual loop without the header (that won't get replaced) # and derive some new IR for this set of blocks this_loop = loop_info.loop this_loop_body = this_loop.body - \ set([this_loop.header]) loop_blocks = { x: func_ir.blocks[x] for x in this_loop_body} new_ir = func_ir.derive(loop_blocks) # Work out what is live on entry and exit so as to prevent # replacement (defined vars can escape, used vars live at the header # need to remain as-is so their references are correct, they can # also escape). usedefs = compute_use_defs(func_ir.blocks) idx = this_loop.header keep = set() keep |= usedefs.usemap[idx] | usedefs.defmap[idx] keep |= func_ir.variable_lifetime.livemap[idx] dont_replace = [x for x in (keep)] # compute the unrolled body unrolled_body = self.inject_loop_body( branches, new_ir, max(func_ir.blocks.keys()) + 1, dont_replace, switch_data) # 6. Patch in the unrolled body and fix up blks = state.func_ir.blocks the_scope = next(iter(blks.values())).scope orig_lbl = tuple(this_loop_body) replace, *delete = orig_lbl unroll, header_block = unrolled_body, this_loop.header unroll_lbl = [x for x in sorted(unroll.blocks.keys())] blks[replace] = transfer_scope(unroll.blocks[unroll_lbl[0]], the_scope) [blks.pop(d) for d in delete] for k in unroll_lbl[1:]: blks[k] = transfer_scope(unroll.blocks[k], the_scope) # stitch up the loop predicate true -> new loop body jump blks[header_block].body[-1].truebr = replace def run_pass(self, state): mutated = False func_ir = state.func_ir # first limit the work by squashing the CFG if possible func_ir.blocks = simplify_CFG(func_ir.blocks) if self._DEBUG: print("-" * 80 + "PASS ENTRY") func_ir.dump() print("-" * 80) # limitations: # 1. No nested unrolls # 2. Opt in via `numba.literal_unroll` # 3. No multiple mix-tuple use # keep running the transform loop until it reports no more changes while(True): stat = self.apply_transform(state) mutated |= stat if not stat: break # reset type inference now we are done with the partial results state.typemap = {} state.calltypes = None return mutated @register_pass(mutates_CFG=True, analysis_only=False) class IterLoopCanonicalization(FunctionPass): """ Transforms loops that are induced by `getiter` into range() driven loops If the typemap is available this will only impact Tuple and UniTuple, if it is not available it will impact all matching loops. """ _name = "iter_loop_canonicalisation" _DEBUG = False # if partial typing info is available it will only look at these types _accepted_types = (types.BaseTuple, types.LiteralList) _accepted_calls = (literal_unroll,) def __init__(self): FunctionPass.__init__(self) def assess_loop(self, loop, func_ir, partial_typemap=None): # it's a iter loop if: # - loop header is driven by an iternext # - the iternext value is a phi derived from getiter() # check header iternexts = [_ for _ in func_ir.blocks[loop.header].find_exprs('iternext')] if len(iternexts) != 1: return False for iternext in iternexts: phi = guard(get_definition, func_ir, iternext.value) if phi is None: return False if getattr(phi, 'op', False) == 'getiter': if partial_typemap: # check that the call site is accepted, until we're # confident that tuple unrolling is behaving require opt-in # guard of `literal_unroll`, remove this later! phi_val_defn = guard(get_definition, func_ir, phi.value) if not isinstance(phi_val_defn, ir.Expr): return False if not phi_val_defn.op == "call": return False call = guard(get_definition, func_ir, phi_val_defn) if call is None or len(call.args) != 1: return False func_var = guard(get_definition, func_ir, call.func) func = guard(get_definition, func_ir, func_var) if func is None or not isinstance(func, (ir.Global, ir.FreeVar)): return False if (func.value is None or func.value not in self._accepted_calls): return False # now check the type is supported ty = partial_typemap.get(call.args[0].name, None) if ty and isinstance(ty, self._accepted_types): return len(loop.entries) == 1 else: return len(loop.entries) == 1 def transform(self, loop, func_ir, cfg): def get_range(a): return range(len(a)) def tokenise(x): return mk_unique_var("CANONICALISER_%s" % x) iternext = [_ for _ in func_ir.blocks[loop.header].find_exprs('iternext')][0] LOC = func_ir.blocks[loop.header].loc get_range_var = ir.Var(func_ir.blocks[loop.header].scope, tokenise('get_range_gbl'), LOC) get_range_global = ir.Global('get_range', get_range, LOC) assgn = ir.Assign(get_range_global, get_range_var, LOC) loop_entry = tuple(loop.entries)[0] entry_block = func_ir.blocks[loop_entry] entry_block.body.insert(0, assgn) iterarg = guard(get_definition, func_ir, iternext.value) if iterarg is not None: iterarg = iterarg.value # look for iternext idx = 0 for stmt in entry_block.body: if isinstance(stmt, ir.Assign): if isinstance(stmt.value, ir.Expr) and stmt.value.op == 'getiter': break idx += 1 else: raise ValueError("problem") # create a range(len(tup)) and inject it call_get_range_var = ir.Var(entry_block.scope, tokenise('call_get_range'), LOC) make_call = ir.Expr.call(get_range_var, (stmt.value.value,), (), LOC) assgn_call = ir.Assign(make_call, call_get_range_var, LOC) entry_block.body.insert(idx, assgn_call) entry_block.body[idx + 1].value.value = call_get_range_var glbls = copy(func_ir.func_id.func.__globals__) from numba.core.inline_closurecall import inline_closure_call inline_closure_call(func_ir, glbls, entry_block, idx, get_range,) kill = entry_block.body.index(assgn) entry_block.body.pop(kill) # find the induction variable + references in the loop header # fixed point iter to do this, it's a bit clunky induction_vars = set() header_block = func_ir.blocks[loop.header] # find induction var ind = [x for x in header_block.find_exprs('pair_first')] for x in ind: induction_vars.add(func_ir.get_assignee(x, loop.header)) # find aliases of the induction var tmp = set() for x in induction_vars: try: # there's not always an alias, e.g. loop from inlined closure tmp.add(func_ir.get_assignee(x, loop.header)) except ValueError: pass induction_vars |= tmp induction_var_names = set([x.name for x in induction_vars]) # Find the downstream blocks that might reference the induction var succ = set() for lbl in loop.exits: succ |= set([x[0] for x in cfg.successors(lbl)]) check_blocks = (loop.body | loop.exits | succ) ^ {loop.header} # replace RHS use of induction var with getitem for lbl in check_blocks: for stmt in func_ir.blocks[lbl].body: if isinstance(stmt, ir.Assign): # check for aliases try: lookup = getattr(stmt.value, 'name', None) except KeyError: continue if lookup and lookup in induction_var_names: stmt.value = ir.Expr.getitem( iterarg, stmt.value, stmt.loc) post_proc = postproc.PostProcessor(func_ir) post_proc.run() def run_pass(self, state): func_ir = state.func_ir cfg = compute_cfg_from_blocks(func_ir.blocks) loops = cfg.loops() mutated = False for header, loop in loops.items(): stat = self.assess_loop(loop, func_ir, state.typemap) if stat: if self._DEBUG: print("Canonicalising loop", loop) self.transform(loop, func_ir, cfg) mutated = True else: if self._DEBUG: print("NOT Canonicalising loop", loop) func_ir.blocks = simplify_CFG(func_ir.blocks) return mutated @register_pass(mutates_CFG=False, analysis_only=False) class PropagateLiterals(FunctionPass): """Implement literal propagation based on partial type inference""" _name = "PropagateLiterals" def __init__(self): FunctionPass.__init__(self) def get_analysis_usage(self, AU): AU.add_required(ReconstructSSA) def run_pass(self, state): func_ir = state.func_ir typemap = state.typemap flags = state.flags if not hasattr(func_ir, '_definitions') \ and not flags.enable_ssa: func_ir._definitions = build_definitions(func_ir.blocks) changed = False for block in func_ir.blocks.values(): for assign in block.find_insts(ir.Assign): value = assign.value if isinstance(value, (ir.Arg, ir.Const, ir.FreeVar, ir.Global)): continue # 1) Don't change return stmt in the form # $return_xyz = cast(value=ABC) # 2) Don't propagate literal values that are not primitives if isinstance(value, ir.Expr) and \ value.op in ('cast', 'build_map', 'build_list', 'build_tuple', 'build_set'): continue target = assign.target if not flags.enable_ssa: # SSA is disabled when doing inlining if guard(get_definition, func_ir, target.name) is None: # noqa: E501 continue # Numba cannot safely determine if an isinstance call # with a PHI node is True/False. For instance, in # the case below, the partial type inference step can coerce # '$z' to float, so any call to 'isinstance(z, int)' would fail. # # def fn(x): # if x > 4: # z = 1 # else: # z = 3.14 # if isinstance(z, int): # print('int') # else: # print('float') # # At the moment, one avoid propagating the literal # value if the argument is a PHI node if isinstance(value, ir.Expr) and value.op == 'call': fn = guard(get_definition, func_ir, value.func.name) if fn is None: continue if not (isinstance(fn, ir.Global) and fn.name == 'isinstance'): # noqa: E501 continue for arg in value.args: # check if any of the args to isinstance is a PHI node iv = func_ir._definitions[arg.name] assert len(iv) == 1 # SSA! if isinstance(iv[0], ir.Expr) and iv[0].op == 'phi': msg = ('isinstance() cannot determine the ' f'type of variable "{arg.unversioned_name}" ' 'due to a branch.') raise errors.NumbaTypeError(msg, loc=assign.loc) lit = typemap.get(target.name, None) if lit and isinstance(lit, types.Literal): # replace assign instruction by ir.Const(lit) iff # lit is a literal value rhs = ir.Const(lit.literal_value, assign.loc) new_assign = ir.Assign(rhs, target, assign.loc) # replace instruction block.insert_after(new_assign, assign) block.remove(assign) changed = True # reset type inference now we are done with the partial results state.typemap = None state.calltypes = None if changed: # Rebuild definitions func_ir._definitions = build_definitions(func_ir.blocks) return changed @register_pass(mutates_CFG=True, analysis_only=False) class LiteralPropagationSubPipelinePass(FunctionPass): """Implement literal propagation based on partial type inference""" _name = "LiteralPropagation" def __init__(self): FunctionPass.__init__(self) def run_pass(self, state): # Determine whether to even attempt this pass... if there's no # `isinstance` as a global or as a freevar then just skip. found = False func_ir = state.func_ir for blk in func_ir.blocks.values(): for asgn in blk.find_insts(ir.Assign): if isinstance(asgn.value, (ir.Global, ir.FreeVar)): if asgn.value.value is isinstance: found = True break if found: break if not found: return False # run as subpipeline from numba.core.compiler_machinery import PassManager from numba.core.typed_passes import PartialTypeInference pm = PassManager("literal_propagation_subpipeline") pm.add_pass(PartialTypeInference, "performs partial type inference") pm.add_pass(PropagateLiterals, "performs propagation of literal values") # rewrite consts / dead branch pruning pm.add_pass(RewriteSemanticConstants, "rewrite semantic constants") pm.add_pass(DeadBranchPrune, "dead branch pruning") pm.finalize() pm.run(state) return True def get_analysis_usage(self, AU): AU.add_required(ReconstructSSA) @register_pass(mutates_CFG=True, analysis_only=False) class LiteralUnroll(FunctionPass): """Implement the literal_unroll semantics""" _name = "literal_unroll" def __init__(self): FunctionPass.__init__(self) def run_pass(self, state): # Determine whether to even attempt this pass... if there's no # `literal_unroll` as a global or as a freevar then just skip. found = False func_ir = state.func_ir for blk in func_ir.blocks.values(): for asgn in blk.find_insts(ir.Assign): if isinstance(asgn.value, (ir.Global, ir.FreeVar)): if asgn.value.value is literal_unroll: found = True break if found: break if not found: return False # run as subpipeline from numba.core.compiler_machinery import PassManager from numba.core.typed_passes import PartialTypeInference pm = PassManager("literal_unroll_subpipeline") # get types where possible to help with list->tuple change pm.add_pass(PartialTypeInference, "performs partial type inference") # make const lists tuples pm.add_pass(TransformLiteralUnrollConstListToTuple, "switch const list for tuples") # recompute partial typemap following IR change pm.add_pass(PartialTypeInference, "performs partial type inference") # canonicalise loops pm.add_pass(IterLoopCanonicalization, "switch iter loops for range driven loops") # rewrite consts pm.add_pass(RewriteSemanticConstants, "rewrite semantic constants") # do the unroll pm.add_pass(MixedContainerUnroller, "performs mixed container unroll") # rewrite dynamic getitem to static getitem as it's possible some more # getitems will now be statically resolvable pm.add_pass(GenericRewrites, "Generic Rewrites") pm.add_pass(RewriteSemanticConstants, "rewrite semantic constants") pm.finalize() pm.run(state) return True @register_pass(mutates_CFG=True, analysis_only=False) class SimplifyCFG(FunctionPass): """Perform CFG simplification""" _name = "simplify_cfg" def __init__(self): FunctionPass.__init__(self) def run_pass(self, state): blks = state.func_ir.blocks new_blks = simplify_CFG(blks) state.func_ir.blocks = new_blks mutated = blks != new_blks return mutated @register_pass(mutates_CFG=False, analysis_only=False) class ReconstructSSA(FunctionPass): """Perform SSA-reconstruction Produces minimal SSA. """ _name = "reconstruct_ssa" def __init__(self): FunctionPass.__init__(self) def run_pass(self, state): state.func_ir = reconstruct_ssa(state.func_ir) self._patch_locals(state) # Rebuild definitions state.func_ir._definitions = build_definitions(state.func_ir.blocks) # Rerun postprocessor to update metadata # example generator_info post_proc = postproc.PostProcessor(state.func_ir) post_proc.run(emit_dels=False) if config.DEBUG or config.DUMP_SSA: name = state.func_ir.func_id.func_qualname print(f"SSA IR DUMP: {name}".center(80, "-")) state.func_ir.dump() return True # XXX detect if it actually got changed def _patch_locals(self, state): # Fix dispatcher locals dictionary type annotation locals_dict = state.get('locals') if locals_dict is None: return first_blk, *_ = state.func_ir.blocks.values() scope = first_blk.scope for parent, redefs in scope.var_redefinitions.items(): if parent in locals_dict: typ = locals_dict[parent] for derived in redefs: locals_dict[derived] = typ