Write prologs and epilogs during CFG construction

kernel/modelverse_jit/cfg_ir.py

@@ -266,11 +266,38 @@ class JitFunctionCall(Value):
             self.target.ref_str(),
             ', '.join(['%s=%s' % (key, val.ref_str()) for key, val in self.argument_list]))
 
+class AllocateRootNode(Value):
+    """A value that produces a new root node. Typically used in function prologs."""
+    def __init__(self):
+        Value.__init__(self)
+
+    def get_dependencies(self):
+        """Gets all definitions and instructions on which this instruction depends."""
+        return []
+
+    def __str__(self):
+        return 'alloc-root-node'
+
+class DeallocateRootNode(Value):
+    """A value that deallocates a root node. Typically used in function epilogs."""
+    def __init__(self, root_node):
+        Value.__init__(self)
+        assert isinstance(root_node, Definition)
+        self.root_node = root_node
+
+    def get_dependencies(self):
+        """Gets all definitions and instructions on which this instruction depends."""
+        return []
+
+    def __str__(self):
+        return 'free-root-node %s' % self.root_node.ref_str()
+
 class DeclareLocal(Value):
     """A value that declares a local variable."""
-    def __init__(self, variable):
+    def __init__(self, variable, root_node):
         Value.__init__(self)
         self.variable = variable
+        self.root_node = root_node
 
     def get_dependencies(self):
         """Gets all definitions and instructions on which this instruction depends."""
@@ -285,7 +312,7 @@ class DeclareLocal(Value):
         return True
 
     def __str__(self):
-        return 'declare-local %d' % self.variable.node_id
+        return 'declare-local %s, %s' % (self.variable, self.root_node.ref_str())
 
 class DeclareGlobal(Value):
     """A value that declares a global variable."""
@@ -319,7 +346,7 @@ class CheckLocalExists(Value):
         return []
 
     def __str__(self):
-        return 'check-local-exists %d' % self.variable.node_id
+        return 'check-local-exists %s' % self.variable
 
 class ResolveLocal(Value):
     """A value that resolves a local as a pointer."""
@@ -332,7 +359,7 @@ class ResolveLocal(Value):
         return []
 
     def __str__(self):
-        return 'resolve-local %d' % self.variable.node_id
+        return 'resolve-local %s' % self.variable
 
 class ResolveGlobal(Value):
     """A value that resolves a global as a pointer."""

kernel/modelverse_jit/cfg_to_tree.py

@@ -4,6 +4,7 @@ import modelverse_jit.cfg_ir as cfg_ir
 import modelverse_jit.cfg_optimization as cfg_optimization
 import modelverse_jit.tree_ir as tree_ir
 import modelverse_jit.runtime as jit_runtime
+import modelverse_jit.bytecode_to_tree as bytecode_to_tree
 
 # The CFG deconstruction code here is based on the relooper algorithm
 # as detailed in https://github.com/kripken/Relooper/blob/master/paper.pdf
@@ -281,12 +282,30 @@ class LoweringState(object):
         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 = {}
+
+    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'
+
+    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)
+
+        if root_node in self.root_edge_names:
+            return self.root_edge_names[root_node]
+
+        result = 'jit_locals%d' % len(self.root_edge_names)
+        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, cfg_ir.Literal):
-                return tree_ir.LiteralInstruction(value.literal)
+                return self.lower_literal(value)
             else:
                 return tree_ir.LoadLocalInstruction(None)
         else:
@@ -312,16 +331,63 @@ class LoweringState(object):
     def lower_definition(self, definition):
         """Lowers the given definition to a tree."""
         instruction = definition.value
-        tree_instruction = self.lower_instruction(instruction)
+        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
 
-    def lower_instruction(self, instruction):
+    def lower_value(self, value):
         """Lowers the given instruction to a tree."""
-        raise NotImplementedError()
+        value_type = type(value)
+        if value_type in LoweringState.value_lowerings:
+            return LoweringState.value_lowerings[value_type](self, value)
+        else:
+            raise jit_runtime.JitCompilationFailedException(
+                "Unknown CFG instruction: '%s'" % value)
+
+    def lower_literal(self, value):
+        """Lowers the given literal value."""
+        return tree_ir.LiteralInstruction(value.literal)
+
+    def lower_check_local_exists(self, value):
+        """Lowers a 'check-value-exists' value."""
+        return tree_ir.LocalExistsInstruction(
+            self.local_name_map.get_local_name(value.variable.node_id))
+
+    def lower_declare_local(self, value):
+        """Lowers a 'declare-local' value."""
+        local_name = self.local_name_map.get_local_name(value.variable.node_id)
+        return tree_ir.create_block(
+            tree_ir.create_new_local_node(
+                local_name,
+                self.load_definition(value.root_node)),
+            tree_ir.LoadLocalInstruction(local_name))
+
+    def lower_alloc_root_node(self, value):
+        """Lowers an 'alloc-root-node' 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,
+                tree_ir.LoadIndexInstruction(
+                    tree_ir.LoadLocalInstruction(jit_runtime.KWARGS_PARAMETER_NAME),
+                    tree_ir.LiteralInstruction('task_root')),
+                self.__get_root_edge_name(value)),
+            tree_ir.LoadLocalInstruction(local_name))
+
+    def lower_free_root_node(self, value):
+        """Lowers a 'free-root-node' value."""
+        return tree_ir.DeleteEdgeInstruction(self.__get_root_edge_name(value.root_node))
+
+    value_lowerings = {
+        cfg_ir.Literal : lower_literal,
+        cfg_ir.CheckLocalExists : lower_check_local_exists,
+        cfg_ir.DeclareLocal : lower_declare_local,
+        cfg_ir.AllocateRootNode : lower_alloc_root_node,
+        cfg_ir.DeallocateRootNode : lower_free_root_node
+    }
 
     def lower_flow(self, flow):
         """Lowers the given (relooped) flow instruction to a tree."""