Fix inliners to run all passes on IR and clean up correctly.

numba/core/compiler.py

@@ -435,11 +435,68 @@ class DefaultPassBuilder(object):
       - nopython
       - objectmode
       - interpreted
+      - typed
+      - untyped
+      - nopython lowering
     """
     @staticmethod
     def define_nopython_pipeline(state, name='nopython'):
         """Returns an nopython mode pipeline based PassManager
         """
+        # compose pipeline from untyped, typed and lowering parts
+        dpb = DefaultPassBuilder
+        pm = PassManager(name)
+        untyped_passes = dpb.define_untyped_pipeline(state)
+        pm.passes.extend(untyped_passes.passes)
+
+        typed_passes = dpb.define_typed_pipeline(state)
+        pm.passes.extend(typed_passes.passes)
+
+        lowering_passes = dpb.define_nopython_lowering_pipeline(state)
+        pm.passes.extend(lowering_passes.passes)
+
+        pm.finalize()
+        return pm
+
+    @staticmethod
+    def define_nopython_lowering_pipeline(state, name='nopython_lowering'):
+        pm = PassManager(name)
+        # legalise
+        pm.add_pass(IRLegalization,
+                    "ensure IR is legal prior to lowering")
+
+        # lower
+        pm.add_pass(NoPythonBackend, "nopython mode backend")
+        pm.add_pass(DumpParforDiagnostics, "dump parfor diagnostics")
+        pm.finalize()
+        return pm
+
+    @staticmethod
+    def define_typed_pipeline(state, name="typed"):
+        """Returns the typed part of the nopython pipeline"""
+        pm = PassManager(name)
+        # typing
+        pm.add_pass(NopythonTypeInference, "nopython frontend")
+        pm.add_pass(AnnotateTypes, "annotate types")
+
+        # strip phis
+        pm.add_pass(PreLowerStripPhis, "remove phis nodes")
+
+        # optimisation
+        pm.add_pass(InlineOverloads, "inline overloaded functions")
+        if state.flags.auto_parallel.enabled:
+            pm.add_pass(PreParforPass, "Preprocessing for parfors")
+        if not state.flags.no_rewrites:
+            pm.add_pass(NopythonRewrites, "nopython rewrites")
+        if state.flags.auto_parallel.enabled:
+            pm.add_pass(ParforPass, "convert to parfors")
+
+        pm.finalize()
+        return pm
+
+    @staticmethod
+    def define_untyped_pipeline(state, name='untyped'):
+        """Returns an untyped part of the nopython pipeline"""
         pm = PassManager(name)
         if state.func_ir is None:
             pm.add_pass(TranslateByteCode, "analyzing bytecode")
@@ -471,29 +528,7 @@ def define_nopython_pipeline(state, name='nopython'):
 
         if state.flags.enable_ssa:
             pm.add_pass(ReconstructSSA, "ssa")
-        # typing
-        pm.add_pass(NopythonTypeInference, "nopython frontend")
-        pm.add_pass(AnnotateTypes, "annotate types")
 
-        # strip phis
-        pm.add_pass(PreLowerStripPhis, "remove phis nodes")
-
-        # optimisation
-        pm.add_pass(InlineOverloads, "inline overloaded functions")
-        if state.flags.auto_parallel.enabled:
-            pm.add_pass(PreParforPass, "Preprocessing for parfors")
-        if not state.flags.no_rewrites:
-            pm.add_pass(NopythonRewrites, "nopython rewrites")
-        if state.flags.auto_parallel.enabled:
-            pm.add_pass(ParforPass, "convert to parfors")
-
-        # legalise
-        pm.add_pass(IRLegalization,
-                    "ensure IR is legal prior to lowering")
-
-        # lower
-        pm.add_pass(NoPythonBackend, "nopython mode backend")
-        pm.add_pass(DumpParforDiagnostics, "dump parfor diagnostics")
         pm.finalize()
         return pm
 

numba/core/inline_closurecall.py

@@ -122,6 +122,13 @@ def run(self):
                 _fix_nested_array(self.func_ir)
 
         if modified:
+            # clean up now dead/unreachable blocks, e.g. unconditionally raising
+            # an exception in an inlined function would render some parts of the
+            # inliner unreachable
+            cfg = compute_cfg_from_blocks(self.func_ir.blocks)
+            for dead in cfg.dead_nodes():
+                del self.func_ir.blocks[dead]
+
             # run dead code elimination
             dead_code_elimination(self.func_ir)
             # do label renaming
@@ -262,6 +269,247 @@ def check_reduce_func(func_ir, func_var):
     return reduce_func
 
 
+class InlineWorker(object):
+    """ A worker class for inlining, this is a more advanced version of
+    `inline_closure_call` in that it permits inlining from function type, Numba
+    IR and code object. It also, runs the entire untyped compiler pipeline on
+    the inlinee to ensure that it is transformed as though it were compiled
+    directly.
+    """
+
+    def __init__(self,
+                 typingctx=None,
+                 targetctx=None,
+                 locals=None,
+                 pipeline=None,
+                 flags=None,
+                 validator=callee_ir_validator,
+                 typemap=None,
+                 calltypes=None):
+        """
+        Instantiate a new InlineWorker, all arguments are optional though some
+        must be supplied together for certain use cases. The methods will refuse
+        to run if the object isn't configured in the manner needed. Args are the
+        same as those in a numba.core.Compiler.state, except the validator which
+        is a function taking Numba IR and validating it for use when inlining
+        (this is optional and really to just provide better error messages about
+        things which the inliner cannot handle like yield in closure).
+        """
+        def check(arg, name):
+            if arg is None:
+                raise TypeError("{} must not be None".format(name))
+
+        from numba.core.compiler import DefaultPassBuilder
+
+        # check the stuff needed to run the more advanced compilation pipeline
+        # is valid if any of it is provided
+        compiler_args = (targetctx, locals, pipeline, flags)
+        compiler_group = [x is not None for x in compiler_args]
+        if any(compiler_group) and not all(compiler_group):
+            check(targetctx, 'targetctx')
+            check(locals, 'locals')
+            check(pipeline, 'pipeline')
+            check(flags, 'flags')
+        elif all(compiler_group):
+            check(typingctx, 'typingctx')
+
+        self._compiler_pipeline = DefaultPassBuilder.define_untyped_pipeline
+
+        self.typingctx = typingctx
+        self.targetctx = targetctx
+        self.locals = locals
+        self.pipeline = pipeline
+        self.flags = flags
+        self.validator = validator
+        self.debug_print = _make_debug_print("InlineWorker")
+
+        # check whether this inliner can also support typemap and calltypes
+        # update and if what's provided is valid
+        pair = (typemap, calltypes)
+        pair_is_none = [x is None for x in pair]
+        if any(pair_is_none) and not all(pair_is_none):
+            msg = ("typemap and calltypes must both be either None or have a "
+                   "value, got: %s, %s")
+            raise TypeError(msg % pair)
+        self._permit_update_type_and_call_maps = not all(pair_is_none)
+        self.typemap = typemap
+        self.calltypes = calltypes
+
+
+    def inline_ir(self, caller_ir, block, i, callee_ir, callee_freevars,
+                  arg_typs=None):
+        """ Inlines the callee_ir in the caller_ir at statement index i of block
+        `block`, callee_freevars are the free variables for the callee_ir. If
+        the callee_ir is derived from a function `func` then this is
+        `func.__code__.co_freevars`. If `arg_typs` is given and the InlineWorker
+        instance was initialized with a typemap and calltypes then they will be
+        appropriately updated based on the arg_typs.
+        """
+        # check that the contents of the callee IR is something that can be
+        # inlined if a validator is present
+        if self.validator is not None:
+            self.validator(callee_ir)
+
+        # save an unmutated copy of the callee_ir to return
+        callee_ir_original = callee_ir.copy()
+        scope = block.scope
+        instr = block.body[i]
+        call_expr = instr.value
+        callee_blocks = callee_ir.blocks
+
+        # 1. relabel callee_ir by adding an offset
+        max_label = max(ir_utils._max_label, max(caller_ir.blocks.keys()))
+        callee_blocks = add_offset_to_labels(callee_blocks, max_label + 1)
+        callee_blocks = simplify_CFG(callee_blocks)
+        callee_ir.blocks = callee_blocks
+        min_label = min(callee_blocks.keys())
+        max_label = max(callee_blocks.keys())
+        #    reset globals in ir_utils before we use it
+        ir_utils._max_label = max_label
+        self.debug_print("After relabel")
+        _debug_dump(callee_ir)
+
+        # 2. rename all local variables in callee_ir with new locals created in
+        # caller_ir
+        callee_scopes = _get_all_scopes(callee_blocks)
+        self.debug_print("callee_scopes = ", callee_scopes)
+        #    one function should only have one local scope
+        assert(len(callee_scopes) == 1)
+        callee_scope = callee_scopes[0]
+        var_dict = {}
+        for var in callee_scope.localvars._con.values():
+            if not (var.name in callee_freevars):
+                new_var = scope.redefine(mk_unique_var(var.name), loc=var.loc)
+                var_dict[var.name] = new_var
+        self.debug_print("var_dict = ", var_dict)
+        replace_vars(callee_blocks, var_dict)
+        self.debug_print("After local var rename")
+        _debug_dump(callee_ir)
+
+        # 3. replace formal parameters with actual arguments
+        callee_func = callee_ir.func_id.func
+        args = _get_callee_args(call_expr, callee_func, block.body[i].loc,
+                                caller_ir)
+
+        # 4. Update typemap
+        if self._permit_update_type_and_call_maps:
+            if arg_typs is None:
+                raise TypeError('arg_typs should have a value not None')
+            self.update_type_and_call_maps(callee_ir, arg_typs)
+
+        self.debug_print("After arguments rename: ")
+        _debug_dump(callee_ir)
+
+        _replace_args_with(callee_blocks, args)
+        # 5. split caller blocks into two
+        new_blocks = []
+        new_block = ir.Block(scope, block.loc)
+        new_block.body = block.body[i + 1:]
+        new_label = next_label()
+        caller_ir.blocks[new_label] = new_block
+        new_blocks.append((new_label, new_block))
+        block.body = block.body[:i]
+        block.body.append(ir.Jump(min_label, instr.loc))
+
+        # 6. replace Return with assignment to LHS
+        topo_order = find_topo_order(callee_blocks)
+        _replace_returns(callee_blocks, instr.target, new_label)
+
+        # remove the old definition of instr.target too
+        if (instr.target.name in caller_ir._definitions
+                and call_expr in caller_ir._definitions[instr.target.name]):
+            # NOTE: target can have multiple definitions due to control flow
+            caller_ir._definitions[instr.target.name].remove(call_expr)
+
+        # 7. insert all new blocks, and add back definitions
+        for label in topo_order:
+            # block scope must point to parent's
+            block = callee_blocks[label]
+            block.scope = scope
+            _add_definitions(caller_ir, block)
+            caller_ir.blocks[label] = block
+            new_blocks.append((label, block))
+        self.debug_print("After merge in")
+        _debug_dump(caller_ir)
+
+        return callee_ir_original, callee_blocks, var_dict, new_blocks
+
+    def inline_function(self, caller_ir, block, i, function, arg_typs=None):
+        """ Inlines the function in the caller_ir at statement index i of block
+        `block`. If `arg_typs` is given and the InlineWorker instance was
+        initialized with a typemap and calltypes then they will be appropriately
+        updated based on the arg_typs.
+        """
+        callee_ir = self.run_untyped_passes(function)
+        freevars = function.__code__.co_freevars
+        return self.inline_ir(caller_ir, block, i, callee_ir, freevars,
+                              arg_typs=arg_typs)
+
+    def run_untyped_passes(self, func):
+        """
+        Run the compiler frontend's untyped passes over the given Python
+        function, and return the function's canonical Numba IR.
+        """
+        from numba.core.compiler import StateDict, _CompileStatus
+        from numba.core.untyped_passes import ExtractByteCode, WithLifting
+        from numba.core import bytecode
+        from numba.parfors.parfor import ParforDiagnostics
+        state = StateDict()
+        state.func_ir = None
+        state.typingctx = self.typingctx
+        state.targetctx = self.targetctx
+        state.locals = self.locals
+        state.pipeline = self.pipeline
+        state.flags = self.flags
+
+        # Disable SSA transformation, the call site won't be in SSA form and
+        # self.inline_ir depends on this being the case.
+        state.flags.enable_ssa = False
+
+        state.func_id = bytecode.FunctionIdentity.from_function(func)
+
+        state.typemap = None
+        state.calltypes = None
+        state.type_annotation = None
+        state.status = _CompileStatus(False, False)
+        state.return_type = None
+        state.parfor_diagnostics = ParforDiagnostics()
+        state.metadata = {}
+
+        ExtractByteCode().run_pass(state)
+        # This is a lie, just need *some* args for the case where an obj mode
+        # with lift is needed
+        state.args = len(state.bc.func_id.pysig.parameters) * (types.pyobject,)
+
+        pm = self._compiler_pipeline(state)
+
+        pm.finalize()
+        pm.run(state)
+        return state.func_ir
+
+    def update_type_and_call_maps(self, callee_ir, arg_typs):
+        """ Updates the type and call maps based on calling callee_ir with arguments
+        from arg_typs"""
+        if not self._permit_update_type_and_call_maps:
+            msg = ("InlineWorker instance not configured correctly, typemap or "
+                   "calltypes missing in initialization.")
+            raise ValueError(msg)
+        from numba.core import typed_passes
+        # call branch pruning to simplify IR and avoid inference errors
+        callee_ir._definitions = ir_utils.build_definitions(callee_ir.blocks)
+        numba.core.analysis.dead_branch_prune(callee_ir, arg_typs)
+        f_typemap, f_return_type, f_calltypes = typed_passes.type_inference_stage(
+                self.typingctx, callee_ir, arg_typs, None)
+        canonicalize_array_math(callee_ir, f_typemap,
+                                f_calltypes, self.typingctx)
+        # remove argument entries like arg.a from typemap
+        arg_names = [vname for vname in f_typemap if vname.startswith("arg.")]
+        for a in arg_names:
+            f_typemap.pop(a)
+        self.typemap.update(f_typemap)
+        self.calltypes.update(f_calltypes)
+
+
 def inline_closure_call(func_ir, glbls, block, i, callee, typingctx=None,
                         arg_typs=None, typemap=None, calltypes=None,
                         work_list=None, callee_validator=None,

numba/core/ir.py

@@ -725,11 +725,11 @@ def __init__(self, exc_class, exc_args, loc):
 
     def __str__(self):
         if self.exc_class is None:
-            return "raise"
+            return "<static> raise"
         elif self.exc_args is None:
-            return "raise %s" % (self.exc_class,)
+            return "<static> raise %s" % (self.exc_class,)
         else:
-            return "raise %s(%s)" % (self.exc_class,
+            return "<static> raise %s(%s)" % (self.exc_class,
                                      ", ".join(map(repr, self.exc_args)))
 
     def get_targets(self):