Implement instruction scheduling for CFG->tree transformation

kernel/modelverse_jit/cfg_to_tree.py

@@ -1,5 +1,6 @@
 """Lowers CFG-IR to tree-IR."""
 
+from collections import defaultdict
 import modelverse_jit.cfg_ir as cfg_ir
 import modelverse_jit.cfg_optimization as cfg_optimization
 import modelverse_jit.tree_ir as tree_ir
@@ -289,23 +290,296 @@ def reloop_function_body(entry_point):
     """Reloops the control-flow graph defined by the given entry point."""
     return reloop_trivial(create_graph_component(entry_point))
 
+def find_inlinable_definitions(entry_point):
+    """Computes a set of definitions which are eligible for inlining, i.e., they
+       are used only once and are consumed in the same basic block in which they
+       are defined."""
+    all_blocks = list(cfg_optimization.get_all_blocks(entry_point))
+
+    # Find all definition users for each definition.
+    def_users = defaultdict(set)
+    def_blocks = {}
+    for block in all_blocks:
+        for parameter_def in block.parameters:
+            def_blocks[parameter_def] = block
+        for definition in block.definitions + [block.flow]:
+            def_blocks[definition] = block
+            for dependency in definition.get_all_dependencies():
+                if not isinstance(dependency, cfg_ir.Branch):
+                    def_users[dependency].add(definition)
+
+    # Find all definitions which are eligible for inlining.
+    eligible_defs = set()
+    for definition, users in def_users.items():
+        if len(users) == 1:
+            single_user = next(iter(users))
+            if def_blocks[single_user] == def_blocks[definition]:
+                eligible_defs.add(definition)
+
+    return eligible_defs
+
+class DefinitionScheduler(object):
+    """Schedules definitions within a basic block."""
+    # We now have the opportunity to schedule definitions in a somewhat intelligent way.
+    # We ought to make sure that we don't reorder operations in an observable way, though.
+    # Specifically:
+    #
+    #     We are allowed to reorder operations which have no side-effects,
+    #     as long as they do not depend on each other. So this is fine:
+    #
+    #         $a = load $x
+    #         $b = load $y
+    #         $c = call-indirect 'jit' f($b, $a)
+    #
+    #     ==>
+    #
+    #         $c = call-indirect 'jit' f(load $y, load $x)
+    #
+    #     But we are not allowed to reorder operations which do have side-effects.
+    #     This is _not_ okay:
+    #
+    #         $a = load $x
+    #         $b = call-indirect 'jit' g()
+    #         $c = store $x, $a
+    #
+    #     ==>
+    #
+    #         $b = call-indirect 'jit' g()
+    #         $c = store $x, load $x
+    #
+    # We can combine the concepts of depending on the results of other instructions
+    # and depending on the potential side-effects of instructions by extending the notion
+    # of what a dependency is: every instruction depends on its explicit dependencies,
+    # and has an implicit dependency on all effectful instructions that precede it.
+    # Additionally, every instruction with side-effects also depends on every instruction
+    # that precedes it, regardless of whether those instructions have side-effects or not.
+    #
+    # The advantage of this way of looking at things is that we can now construct a dependency
+    # graph. For example, this is the dependency graph for the second example.
+    #
+    #                                $c = store $x, $a
+    #                               /                 \
+    #                              /                   \
+    #                             v                     v
+    #                    $a = load $x <------------- $b = call-indirect 'jit' g()
+    #
+    # Any topological sort of the dependency graph is a valid instruction scheduling, but we
+    # specifically want to grow a tree that uses few temporaries.
+    #
+    # We can accomplish this by growing the tree in reverse: we pick a definition on which no other
+    # definition depends, and remove it from the dependency graph. We then iterate over the
+    # definition's dependencies in reverse. If a dependency is encountered that is eligible for
+    # inlining, i.e., it is used only once and is defined in the basic block where it is used,
+    # then we really want to schedule it inline -- doing so allows us to avoid defining a temporary.
+    # Here's the algorithm that we'll use to schedule dependencies.
+    #
+    # while (dependency has not been scheduled) and (another definition depends on the dependency):
+    #     pick a definition that depends on the dependency and schedule it
+    #     store the scheduled definition's value in a temporary
+    #
+    # if (dependency has not been scheduled) and (dependency is eligible for inlining):
+    #     schedule the dependency inline
+    # else if (dependency has not been scheduled):
+    #     schedule the dependency, and store it in a temporary
+    # else:
+    #     load the (already scheduled) dependency's value from a temporary
+    #
+    # Note that it is possible for another definition to also depend on a dependency, despite having
+    # the dependency used only once. This merely means that at least one dependency is based on an
+    # instruction's side-effects.
+    #
+    # Okay, so that's the theory. In this class, we will use the following data structure to
+    # improve performance:
+    #
+    #     * We will implement our graph as a map of definitions to a
+    #       (outgoing edges, incoming edges) tuple. This will allow us to quickly test if a
+    #       definition is a dependency of some other definition, and will also allow us to
+    #       erase definitions from the graph efficiently.
+    #
+    #     * We will not create implicit (side-effect--based) dependencies beyond the last
+    #       definition with side-effects. Said definition already depends on every definition
+    #       that precedes it, so those edges are simply redundant.
+    #
+    def __init__(self, lowering_state, definitions, flow):
+        self.lowering_state = lowering_state
+        self.dependency_graph = {}
+        self.construct_dependency_graph(definitions + [flow])
+
+    def add_dependency(self, definition, dependency):
+        """Makes the given definition dependent on the given dependency in the
+           dependency graph."""
+        def_outgoing_edges, _ = self.dependency_graph[definition]
+        _, dep_incoming_edges = self.dependency_graph[dependency]
+        def_outgoing_edges.add(dependency)
+        dep_incoming_edges.add(definition)
+
+    def construct_dependency_graph(self, definitions):
+        """Construct the dependency graph for the given set of definitions."""
+        def_set = set(definitions)
+
+        # Initialize the dependency graph for every definition in the basic block.
+        for definition in definitions:
+            self.dependency_graph[definition] = (set(), set())
+
+        # Construct the dependency graph.
+        last_def_batch = []
+        last_effectful_def = None
+        for definition in definitions:
+            # Explicit dependencies.
+            for dep in definition.get_all_dependencies():
+                if dep in def_set:
+                    self.add_dependency(definition, dep)
+
+            # Implicit dependency on the previous definition with side-effects.
+            if last_effectful_def is not None:
+                self.add_dependency(definition, last_effectful_def)
+
+            # Implicit dependency of definitions with side-effects on all definitions
+            # that precede them. Implementation detail: we don't actually make
+            # definitions dependent on _all_ definitions that precede them. It suffices
+            # to insert a dependency on every definition between the current definition
+            # and the previous definition with side-effects.
+            if definition.has_side_effects():
+                for implicit_dependency in last_def_batch:
+                    self.add_dependency(definition, implicit_dependency)
+                last_def_batch = []
+                last_effectful_def = definition
+            else:
+                last_def_batch.append(definition)
+
+    def remove_definition(self, definition):
+        """Removes the given definition from the graph."""
+        assert not self.has_dependents(definition)
+
+        # Remove the definition from the graph.
+        outgoing_edges, _ = self.dependency_graph[definition]
+        del self.dependency_graph[definition]
+
+        # Remove the definition's outgoing edges from the graph.
+        for dependency in outgoing_edges:
+            _, incoming_edges = self.dependency_graph[dependency]
+            incoming_edges.remove(definition)
+
+    def has_scheduled(self, definition):
+        """Tests if the given definition has already been scheduled."""
+        return definition not in self.dependency_graph
+
+    def get_schedulable_definition(self):
+        """Finds a schedulable definition. Returns None if the dependency graph is empty."""
+        if len(self.dependency_graph) == 0:
+            return None
+
+        for node in self.dependency_graph.keys():
+            if not self.has_dependents(node):
+                return node
+
+        raise ValueError(
+            'Could not find an instruction to schedule. Dependency graph: %s' %
+            self.dependency_graph)
+
+    def has_dependents(self, definition):
+        """Tests if there is at least one definition in the dependency graph that
+           depends on the given definition."""
+        return len(self.get_dependents(definition)) > 0
+
+    def get_dependents(self, definition):
+        """Gets the set of definitions that depend on the given definition."""
+        _, incoming_edges = self.dependency_graph[definition]
+        return incoming_edges
+
+    def use(self, definition):
+        """Creates a tree that produces the given definition's value."""
+        def_value = cfg_ir.get_def_value(definition)
+        if isinstance(def_value, LoweringState.inline_value_types):
+            return self.lowering_state.lower_value(def_value)
+        elif isinstance(def_value, cfg_ir.AllocateRootNode):
+            return tree_ir.LoadLocalInstruction(
+                self.lowering_state.get_root_node_name(def_value))
+        elif not self.has_scheduled(definition):
+            return self.schedule(definition)
+        elif definition.has_value():
+            return self.lowering_state.create_definition_load(definition)
+        else:
+            return tree_ir.EmptyInstruction()
+
+    def __schedule_definition(self, definition, store_and_forget=False):
+        """Schedules the given definition, excluding definitions that are dependent on it.
+           A list of trees is returned."""
+        statements = []
+        if not self.has_scheduled(definition):
+            self.remove_definition(definition)
+            if isinstance(definition, cfg_ir.FlowInstruction):
+                statements.append(self.lowering_state.lower_value(definition))
+            else:
+                # We're sure that we're dealing with a bona fide cfg_ir.Definition now.
+                definition_value = definition.value
+                if cfg_ir.is_value_def(definition_value, LoweringState.inline_value_types):
+                    pass
+                elif isinstance(definition_value, cfg_ir.Definition):
+                    value_use = self.use(definition_value)
+                    if definition_value.has_value():
+                        if store_and_forget:
+                            def_load = self.lowering_state.create_definition_load(definition)
+                            statements.append(
+                                tree_ir.IgnoreInstruction(def_load.create_store(value_use)))
+                        else:
+                            statements.append(value_use)
+                elif (not store_and_forget
+                      and definition in self.lowering_state.inlinable_definitions):
+                    statements.append(self.lowering_state.lower_value(definition_value))
+                else:
+                    lowered_value = self.lowering_state.lower_value(definition_value)
+                    if definition.has_value():
+                        def_load = self.lowering_state.create_definition_load(definition)
+                        statements.append(
+                            tree_ir.IgnoreInstruction(def_load.create_store(lowered_value)))
+                        if not store_and_forget:
+                            statements.append(def_load)
+                    else:
+                        statements.append(lowered_value)
+        elif (not store_and_forget
+              and definition.has_value()
+              and not cfg_ir.is_value_def(definition, LoweringState.inline_value_types)):
+            statements.append(self.lowering_state.create_definition_load(definition))
+
+        return statements
+
+    def schedule(self, definition, store_and_forget=False):
+        """Schedules the given definition. The resulting tree is returned."""
+        assert not self.has_scheduled(definition)
+
+        statements = []
+        dependents = self.get_dependents(definition)
+        while not self.has_scheduled(definition) and len(dependents) > 0:
+            dependent = next(iter(dependents))
+            dependent_tree = self.schedule(dependent, True)
+            statements.append(dependent_tree)
+
+        statements.reverse()
+        statements = self.__schedule_definition(definition, store_and_forget) + statements
+
+        result = tree_ir.create_block(*statements)
+        return result
+
 class LoweringState(object):
     """Stores information related to the relooper->tree conversion."""
-    def __init__(self, jit):
+    def __init__(self, jit, inlinable_definitions):
         self.jit = jit
         self.label_variable = tree_ir.VariableName('__label')
         self.definition_loads = {}
         self.local_name_map = bytecode_to_tree.LocalNameMap()
         self.root_edge_names = {}
+        self.inlinable_definitions = inlinable_definitions
+        self.scheduler = None
 
     def __get_root_edge_name(self, root_node):
         """Gets the name of the given root edge's variable."""
-        return self.__get_root_node_name(root_node) + '_edge'
+        return self.get_root_node_name(root_node) + '_edge'
 
-    def __get_root_node_name(self, root_node):
+    def get_root_node_name(self, root_node):
         """Gets the name of the given root node's variable."""
         if isinstance(root_node, cfg_ir.Definition):
-            return self.__get_root_edge_name(root_node.value)
+            return self.get_root_node_name(root_node.value)
 
         if root_node in self.root_edge_names:
             return self.root_edge_names[root_node]
@@ -314,55 +588,40 @@ class LoweringState(object):
         self.root_edge_names[root_node] = result
         return result
 
-    def __create_value_load(self, value):
-        """Creates a tree that loads the given value."""
-        if value.has_value():
-            if isinstance(value, LoweringState.inline_value_types):
-                return self.lower_value(value)
-            else:
-                return tree_ir.LoadLocalInstruction(None)
-        else:
-            return tree_ir.LiteralInstruction(None)
-
-    def load_definition(self, definition):
+    def create_definition_load(self, definition):
         """Loads the given definition's variable."""
         if definition in self.definition_loads:
             return self.definition_loads[definition]
 
-        if isinstance(definition.value, cfg_ir.Definition):
-            result = self.load_definition(definition.value)
-        else:
-            result = self.__create_value_load(definition.value)
+        result = tree_ir.LoadLocalInstruction(None)
         self.definition_loads[definition] = result
         return result
 
+    def use_definition(self, definition):
+        """Loads the given definition's value."""
+        return self.scheduler.use(definition)
+
     def lower_block(self, block):
         """Lowers the given (relooped) block to a tree."""
         statements = []
-        for definition in block.definitions:
-            statements.append(self.lower_definition(definition))
-        statements.append(self.lower_flow(block.flow))
-        return tree_ir.create_block(*statements)
+        self.scheduler = DefinitionScheduler(self, block.definitions, block.flow)
+        while 1:
+            schedulable_def = self.scheduler.get_schedulable_definition()
+            if schedulable_def is None:
+                break
 
-    def lower_definition(self, definition):
-        """Lowers the given definition to a tree."""
-        instruction = definition.value
-        if (isinstance(instruction, cfg_ir.Definition)
-                or isinstance(instruction, LoweringState.inline_value_types)):
-            return tree_ir.EmptyInstruction()
+            statements.append(self.scheduler.schedule(schedulable_def))
 
-        tree_instruction = self.lower_value(instruction)
-        def_load = self.load_definition(definition)
-        if isinstance(def_load, tree_ir.LocalInstruction):
-            return def_load.create_store(tree_instruction)
-        else:
-            return tree_instruction
+        self.scheduler = None
+
+        statements.reverse()
+        return tree_ir.create_block(*statements)
 
     def lower_value(self, value):
         """Lowers the given instruction to a tree."""
         value_type = type(value)
-        if value_type in LoweringState.value_lowerings:
-            return LoweringState.value_lowerings[value_type](self, value)
+        if value_type in LoweringState.instruction_lowerings:
+            return LoweringState.instruction_lowerings[value_type](self, value)
         else:
             raise jit_runtime.JitCompilationFailedException(
                 "Unknown CFG instruction: '%s'" % value)
@@ -382,7 +641,7 @@ class LoweringState(object):
         return tree_ir.create_block(
             tree_ir.create_new_local_node(
                 local_name,
-                self.load_definition(value.root_node)),
+                self.use_definition(value.root_node)),
             tree_ir.LoadLocalInstruction(local_name))
 
     def lower_resolve_local(self, value):
@@ -430,7 +689,7 @@ class LoweringState(object):
 
     def lower_alloc_root_node(self, value):
         """Lowers an 'alloc-root-node' value."""
-        local_name = tree_ir.VariableName(self.__get_root_node_name(value))
+        local_name = tree_ir.VariableName(self.get_root_node_name(value))
         return tree_ir.create_block(
             tree_ir.create_new_local_node(
                 local_name,
@@ -447,32 +706,30 @@ class LoweringState(object):
 
     def lower_load_pointer(self, value):
         """Lowers a 'load' value."""
-        return bytecode_to_tree.create_access(self.load_definition(value.pointer))
+        return bytecode_to_tree.create_access(self.use_definition(value.pointer))
 
     def lower_store_pointer(self, value):
         """Lowers a 'store' value."""
-        return bytecode_to_tree.create_assign(
-            self.load_definition(value.pointer), self.load_definition(value.value))
+        ptr, val = self.use_definition(value.pointer), self.use_definition(value.value)
+        return bytecode_to_tree.create_assign(ptr, val)
 
     def lower_read(self, value):
         """Lowers a 'read' value."""
         return tree_ir.ReadValueInstruction(
-            self.load_definition(value.node))
+            self.use_definition(value.node))
 
     def lower_create_node(self, value):
         """Lowers a 'create-node' value."""
         if value.value.has_value():
             return tree_ir.CreateNodeWithValueInstruction(
-                self.load_definition(value.value))
+                self.use_definition(value.value))
         else:
             return tree_ir.CreateNodeInstruction()
 
     def lower_binary(self, value):
         """Lowers a 'binary' value."""
-        return tree_ir.BinaryInstruction(
-            self.load_definition(value.lhs),
-            value.operator,
-            self.load_definition(value.rhs))
+        lhs, rhs = self.use_definition(value.lhs), self.use_definition(value.rhs)
+        return tree_ir.BinaryInstruction(lhs, value.operator, rhs)
 
     def lower_input(self, _):
         """Lowers an 'input' value."""
@@ -480,7 +737,7 @@ class LoweringState(object):
 
     def lower_output(self, value):
         """Lowers an 'output' value."""
-        return bytecode_to_tree.create_output(self.load_definition(value.value))
+        return bytecode_to_tree.create_output(self.use_definition(value.value))
 
     def lower_direct_call(self, value):
         """Lowers a direct function call."""
@@ -494,42 +751,21 @@ class LoweringState(object):
 
     def lower_indirect_call(self, value):
         """Lowers an indirect function call."""
-        return bytecode_to_tree.create_indirect_call(
-            self.load_definition(value.target),
-            [(name, self.load_definition(arg)) for name, arg in value.argument_list])
-
-    inline_value_types = (cfg_ir.Literal, cfg_ir.ResolveLocal, cfg_ir.FunctionParameter)
-
-    value_lowerings = {
-        cfg_ir.Literal : lower_literal,
-        cfg_ir.CheckLocalExists : lower_check_local_exists,
-        cfg_ir.DeclareLocal : lower_declare_local,
-        cfg_ir.ResolveLocal : lower_resolve_local,
-        cfg_ir.DeclareGlobal : lower_declare_global,
-        cfg_ir.ResolveGlobal : lower_resolve_global,
-        cfg_ir.FunctionParameter : lower_function_parameter,
-        cfg_ir.AllocateRootNode : lower_alloc_root_node,
-        cfg_ir.DeallocateRootNode : lower_free_root_node,
-        cfg_ir.LoadPointer : lower_load_pointer,
-        cfg_ir.StoreAtPointer : lower_store_pointer,
-        cfg_ir.Read : lower_read,
-        cfg_ir.CreateNode : lower_create_node,
-        cfg_ir.Binary : lower_binary,
-        cfg_ir.Input : lower_input,
-        cfg_ir.Output : lower_output,
-        cfg_ir.DirectFunctionCall : lower_direct_call,
-        cfg_ir.IndirectFunctionCall : lower_indirect_call
-    }
+        target = self.use_definition(value.target)
+        args = [(name, self.use_definition(arg))
+                for name, arg in value.argument_list]
+        return bytecode_to_tree.create_indirect_call(target, args)
 
     def lower_simple_positional_call(self, value):
         """Lowers a direct call that uses the 'simple-positional' calling convention."""
         return tree_ir.CallInstruction(
             tree_ir.LoadGlobalInstruction(value.target_name),
-            [self.load_definition(arg) for _, arg in value.argument_list])
+            [self.use_definition(arg) for _, arg in value.argument_list])
 
     def lower_jit_call(self, value):
         """Lowers a direct call that uses the 'jit' calling convention."""
-        arg_list = [(name, self.load_definition(arg)) for name, arg in value.argument_list]
+        arg_list = [(name, self.use_definition(arg))
+                    for name, arg in value.argument_list]
         intrinsic = self.jit.get_intrinsic(value.target_name)
         if intrinsic is not None:
             return bytecode_to_tree.apply_intrinsic(intrinsic, arg_list)
@@ -539,20 +775,6 @@ class LoweringState(object):
                 arg_list,
                 tree_ir.LoadLocalInstruction(jit_runtime.KWARGS_PARAMETER_NAME))
 
-    call_lowerings = {
-        cfg_ir.SIMPLE_POSITIONAL_CALLING_CONVENTION : lower_simple_positional_call,
-        cfg_ir.JIT_CALLING_CONVENTION : lower_jit_call
-    }
-
-    def lower_flow(self, flow):
-        """Lowers the given (relooped) flow instruction to a tree."""
-        flow_type = type(flow)
-        if flow_type in LoweringState.flow_lowerings:
-            return LoweringState.flow_lowerings[flow_type](self, flow)
-        else:
-            raise jit_runtime.JitCompilationFailedException(
-                "Unknown CFG flow instruction: '%s'" % flow)
-
     def lower_jump(self, flow):
         """Lowers the given 'jump' flow instruction to a tree."""
         return self.lower_branch(flow.branch)
@@ -560,17 +782,17 @@ class LoweringState(object):
     def lower_select(self, flow):
         """Lowers the given 'select' flow instruction to a tree."""
         return tree_ir.SelectInstruction(
-            self.load_definition(flow.condition),
+            self.use_definition(flow.condition),
             self.lower_branch(flow.if_branch),
             self.lower_branch(flow.else_branch))
 
     def lower_return(self, flow):
         """Lowers the given 'return' flow instruction to a tree."""
-        return tree_ir.ReturnInstruction(self.load_definition(flow.value))
+        return tree_ir.ReturnInstruction(self.use_definition(flow.value))
 
     def lower_throw(self, flow):
         """Lowers the given 'throw' flow instruction to a tree."""
-        return tree_ir.RaiseInstruction(self.load_definition(flow.exception))
+        return tree_ir.RaiseInstruction(self.use_definition(flow.exception))
 
     def lower_unreachable(self, _):
         """Lowers the given 'unreachable' flow instruction to a tree."""
@@ -587,13 +809,34 @@ class LoweringState(object):
     def lower_branch(self, branch):
         """Lowers the given (relooped) branch to a tree."""
         for param, arg in zip(branch.block.parameters, branch.arguments):
-            self.load_definition(param).create_store(self.load_definition(arg))
+            self.create_definition_load(param).create_store(self.use_definition(arg))
+
+        return tree_ir.IgnoreInstruction(
+            tree_ir.StoreLocalInstruction(
+                self.label_variable,
+                tree_ir.LiteralInstruction(branch.block.index)))
 
-        return tree_ir.StoreLocalInstruction(
-            self.label_variable,
-            tree_ir.LiteralInstruction(branch.block.index))
+    inline_value_types = (cfg_ir.Literal, cfg_ir.ResolveLocal, cfg_ir.FunctionParameter)
 
-    flow_lowerings = {
+    instruction_lowerings = {
+        cfg_ir.Literal : lower_literal,
+        cfg_ir.CheckLocalExists : lower_check_local_exists,
+        cfg_ir.DeclareLocal : lower_declare_local,
+        cfg_ir.ResolveLocal : lower_resolve_local,
+        cfg_ir.DeclareGlobal : lower_declare_global,
+        cfg_ir.ResolveGlobal : lower_resolve_global,
+        cfg_ir.FunctionParameter : lower_function_parameter,
+        cfg_ir.AllocateRootNode : lower_alloc_root_node,
+        cfg_ir.DeallocateRootNode : lower_free_root_node,
+        cfg_ir.LoadPointer : lower_load_pointer,
+        cfg_ir.StoreAtPointer : lower_store_pointer,
+        cfg_ir.Read : lower_read,
+        cfg_ir.CreateNode : lower_create_node,
+        cfg_ir.Binary : lower_binary,
+        cfg_ir.Input : lower_input,
+        cfg_ir.Output : lower_output,
+        cfg_ir.DirectFunctionCall : lower_direct_call,
+        cfg_ir.IndirectFunctionCall : lower_indirect_call,
         cfg_ir.JumpFlow : lower_jump,
         cfg_ir.SelectFlow : lower_select,
         cfg_ir.ReturnFlow : lower_return,
@@ -602,3 +845,16 @@ class LoweringState(object):
         BreakFlow : lower_break,
         ContinueFlow : lower_continue
     }
+
+    call_lowerings = {
+        cfg_ir.SIMPLE_POSITIONAL_CALLING_CONVENTION : lower_simple_positional_call,
+        cfg_ir.JIT_CALLING_CONVENTION : lower_jit_call
+    }
+
+def lower_flow_graph(entry_point, jit):
+    """Lowers the control-flow graph defined by the given entry point to tree IR."""
+    cfg_lowerer = LoweringState(jit, find_inlinable_definitions(entry_point))
+    lowered_body = reloop_function_body(entry_point).lower(cfg_lowerer)
+    lowered_body = tree_ir.CompoundInstruction(
+        cfg_lowerer.lower_jump(cfg_ir.create_jump(entry_point)), lowered_body)
+    return tree_ir.map_and_simplify(lambda x: x, lowered_body)

kernel/modelverse_jit/jit.py

@@ -47,6 +47,15 @@ def optimize_tree_ir(instruction):
     """Optimizes an IR tree."""
     return map_and_simplify_generator(expand_constant_read, instruction)
 
+def create_bare_function(function_name, parameter_list, function_body):
+    """Creates a function definition from the given function name, parameter list
+       and function body. No prolog is included."""
+    # Wrap the IR in a function definition, give it a unique name.
+    return tree_ir.DefineFunctionInstruction(
+        function_name,
+        parameter_list + ['**' + jit_runtime.KWARGS_PARAMETER_NAME],
+        function_body)
+
 def create_function(
         function_name, parameter_list, param_dict,
         body_param_dict, function_body):
@@ -81,13 +90,7 @@ def create_function(
     constructed_body = tree_ir.protect_temporaries_from_gc(
         constructed_body, tree_ir.LoadLocalInstruction(jit_runtime.LOCALS_NODE_NAME))
 
-    # Wrap the IR in a function definition, give it a unique name.
-    constructed_function = tree_ir.DefineFunctionInstruction(
-        function_name,
-        parameter_list + ['**' + jit_runtime.KWARGS_PARAMETER_NAME],
-        constructed_body)
-
-    return constructed_function
+    return create_bare_function(function_name, parameter_list, constructed_body)
 
 def print_value(val):
     """A thin wrapper around 'print'."""
@@ -477,12 +480,12 @@ class ModelverseJit(object):
                             str,
                             cfg_optimization.get_all_reachable_blocks(
                                 bytecode_analyzer.entry_point)))))
-            cfg_lowerer = cfg_to_tree.LoweringState(self)
-            lowered_body = cfg_to_tree.reloop_function_body(
-                bytecode_analyzer.entry_point).lower(cfg_lowerer)
+            cfg_func = create_bare_function(
+                function_name, parameter_list,
+                cfg_to_tree.lower_flow_graph(bytecode_analyzer.entry_point, self))
             self.jit_code_log_function(
                 "Lowered CFG for function '%s' at '%d':\n%s" % (
-                    function_name, body_id, lowered_body))
+                    function_name, body_id, cfg_func))
 
         yield [("END_TRY", [])]
         del self.compilation_dependencies[body_id]