Add JIT support for 'output' instructions

kernel/modelverse_jit/jit.py

@@ -1,6 +1,9 @@
 import modelverse_kernel.primitives as primitive_functions
 import modelverse_jit.tree_ir as tree_ir
 
+KWARGS_PARAMETER_NAME = "remainder"
+"""The name of the kwargs parameter in jitted functions."""
+
 class JitCompilationFailedException(Exception):
     """A type of exception that is raised when the jit fails to compile a function."""
     pass
@@ -51,8 +54,12 @@ class ModelverseJit(object):
             while True:
                 inp = yield gen.send(inp)
         except primitive_functions.PrimitiveFinished as e:
-            pass
+            constructed_ir = e.result
+        except JitCompilationFailedException:
+            self.mark_no_jit(body_id)
+            raise
 
+        print(constructed_ir)
         self.mark_no_jit(body_id)
         raise JitCompilationFailedException("Can't JIT function body at " + str(body_id))
 
@@ -169,8 +176,8 @@ class AnalysisState(object):
                     tree_ir.CompoundInstruction(
                         tree_ir.SelectInstruction(
                             tree_ir.ReadValueInstruction(cond_r),
-                            tree_ir.BreakInstruction(),
-                            tree_ir.EmptyInstruction()),
+                            tree_ir.EmptyInstruction(),
+                            tree_ir.BreakInstruction()),
                         body_r)))
 
     def analyze_constant(self, instruction_id):
@@ -179,9 +186,82 @@ class AnalysisState(object):
         raise primitive_functions.PrimitiveFinished(
             tree_ir.LiteralInstruction(node_id))
 
+    def analyze_output(self, instruction_id):
+        """Tries to analyze the given 'output' instruction."""
+
+        # The plan is to basically generate this tree:
+        #
+        # value = <some tree>
+        # last_output, last_output_link, new_last_output = \
+        #                 yield [("RD", [user_root, "last_output"]),
+        #                        ("RDE", [user_root, "last_output"]),
+        #                        ("CN", []),
+        #                       ]
+        # _, _, _, _ = \
+        #                 yield [("CD", [last_output, "value", value]),
+        #                        ("CD", [last_output, "next", new_last_output]),
+        #                        ("CD", [user_root, "last_output", new_last_output]),
+        #                        ("DE", [last_output_link])
+        #                       ]
+
+        value_id, = yield [("RD", [instruction_id, "value"])]
+        gen = self.analyze(value_id)
+        try:
+            inp = None
+            while True:
+                inp = yield gen.send(inp)
+        except primitive_functions.PrimitiveFinished as e:
+            value_local = tree_ir.StoreLocalInstruction('value', e.result)
+
+        store_user_root = tree_ir.StoreLocalInstruction(
+            'user_root',
+            tree_ir.LoadIndexInstruction(
+                tree_ir.LoadLocalInstruction(KWARGS_PARAMETER_NAME),
+                tree_ir.LiteralInstruction('user_root')))
+
+        last_output = tree_ir.StoreLocalInstruction(
+            'last_output',
+            tree_ir.ReadDictionaryValueInstruction(
+                store_user_root.create_load(),
+                tree_ir.LiteralInstruction('last_output')))
+
+        last_output_link = tree_ir.StoreLocalInstruction(
+            'last_output_link',
+            tree_ir.ReadDictionaryEdgeInstruction(
+                store_user_root.create_load(),
+                tree_ir.LiteralInstruction('last_output')))
+
+        new_last_output = tree_ir.StoreLocalInstruction(
+            'new_last_output',
+            tree_ir.CreateNodeInstruction())
+
+        result = tree_ir.create_block(
+            value_local,
+            store_user_root,
+            last_output,
+            last_output_link,
+            new_last_output,
+            tree_ir.CreateDictionaryEdgeInstruction(
+                last_output.create_load(),
+                tree_ir.LiteralInstruction('value'),
+                value_local.create_load()),
+            tree_ir.CreateDictionaryEdgeInstruction(
+                last_output.create_load(),
+                tree_ir.LiteralInstruction('next'),
+                new_last_output.create_load()),
+            tree_ir.CreateDictionaryEdgeInstruction(
+                store_user_root.create_load(),
+                tree_ir.LiteralInstruction('last_output'),
+                new_last_output.create_load()),
+            tree_ir.DeleteEdgeInstruction(last_output_link.create_load()))
+
+        raise primitive_functions.PrimitiveFinished(result)
+
     instruction_analyzers = {
         'if' : analyze_if,
         'while' : analyze_while,
         'return' : analyze_return,
-        'constant' : analyze_constant
+        'constant' : analyze_constant,
+        'output' : analyze_output
     }
+

kernel/modelverse_jit/tree_ir.py

@@ -288,6 +288,71 @@ class StateInstruction(Instruction):
 
         code_generator.append_state_definition(self, self.get_opcode(), args)
 
+class LocalInstruction(Instruction):
+    """A base class for instructions that access local variables."""
+    def __init__(self, name):
+        Instruction.__init__(self)
+        self.name = name
+
+    def create_load(self):
+        """Creates an instruction that loads the variable referenced by this instruction."""
+        return LoadLocalInstruction(self.name)
+
+    def generate_python_use(self, code_generator):
+        """Generates a Python expression that retrieves this instruction's
+           result. The expression is returned as a string."""
+        return self.name
+
+class StoreLocalInstruction(LocalInstruction):
+    """An instruction that stores a value in a local variable."""
+    def __init__(self, name, value):
+        LocalInstruction.__init__(self, name)
+        self.value = value
+
+    def simplify(self):
+        """Applies basic simplification to this instruction and its children."""
+        return StoreLocalInstruction(self.name, self.value.simplify())
+
+    def generate_python_def(self, code_generator):
+        """Generates a Python statement that executes this instruction.
+           The statement is appended immediately to the code generator."""
+
+        if self.value.has_definition():
+            self.value.generate_python_def(code_generator)
+        code_generator.append_line(
+            '%s = %s' % (self.name, self.value.generate_python_use(code_generator)))
+
+class LoadLocalInstruction(LocalInstruction):
+    """An instruction that loads a value from a local variable."""
+    def has_definition(self):
+        """Tells if this instruction requires a definition."""
+        return False
+
+class LoadIndexInstruction(Instruction):
+    """An instruction that produces a value by indexing a specified expression with
+       a given key."""
+    def __init__(self, indexed, key):
+        Instruction.__init__(self)
+        self.indexed = indexed
+        self.key = key
+
+    def has_definition(self):
+        """Tells if this instruction requires a definition."""
+        return False
+
+    def generate_python_use(self, code_generator):
+        """Generates a Python expression that retrieves this instruction's
+           result. The expression is returned as a string."""
+        if self.indexed.has_definition():
+            self.indexed.generate_python_def(code_generator)
+
+        if self.key.has_definition():
+            self.key.generate_python_def(code_generator)
+
+        return "%s[%s]" % (
+            self.indexed.generate_python_use(code_generator),
+            self.key.generate_python_use(code_generator))
+
 class ReadValueInstruction(StateInstruction):
     """An instruction that reads a value from a node."""
 
@@ -307,6 +372,120 @@ class ReadValueInstruction(StateInstruction):
         """Gets this state instruction's argument list."""
         return [self.node_id]
 
+class ReadDictionaryValueInstruction(StateInstruction):
+    """An instruction that reads a dictionary value."""
+
+    def __init__(self, node_id, key):
+        StateInstruction.__init__(self)
+        self.node_id = node_id
+        self.key = key
+
+    def simplify(self):
+        """Applies basic simplification to this instruction and its children."""
+        return ReadDictionaryValueInstruction(
+            self.node_id.simplify(),
+            self.key.simplify())
+
+    def get_opcode(self):
+        """Gets the opcode for this state instruction."""
+        return "RD"
+
+    def get_arguments(self):
+        """Gets this state instruction's argument list."""
+        return [self.node_id, self.key]
+
+class ReadDictionaryEdgeInstruction(StateInstruction):
+    """An instruction that reads a dictionary edge."""
+
+    def __init__(self, node_id, key):
+        StateInstruction.__init__(self)
+        self.node_id = node_id
+        self.key = key
+
+    def simplify(self):
+        """Applies basic simplification to this instruction and its children."""
+        return ReadDictionaryEdgeInstruction(
+            self.node_id.simplify(),
+            self.key.simplify())
+
+    def get_opcode(self):
+        """Gets the opcode for this state instruction."""
+        return "RDE"
+
+    def get_arguments(self):
+        """Gets this state instruction's argument list."""
+        return [self.node_id, self.key]
+
+class CreateNodeInstruction(StateInstruction):
+    """An instruction that creates an empty node."""
+
+    def get_opcode(self):
+        """Gets the opcode for this state instruction."""
+        return "CN"
+
+    def get_arguments(self):
+        """Gets this state instruction's argument list."""
+        return []
+
+class CreateDictionaryEdgeInstruction(StateInstruction):
+    """An instruction that creates a dictionary edge."""
+
+    def __init__(self, source_id, key, target_id):
+        StateInstruction.__init__(self)
+        self.source_id = source_id
+        self.key = key
+        self.target_id = target_id
+
+    def simplify(self):
+        """Applies basic simplification to this instruction and its children."""
+        return CreateDictionaryEdgeInstruction(
+            self.source_id.simplify(),
+            self.key.simplify(),
+            self.target_id.simplify())
+
+    def get_opcode(self):
+        """Gets the opcode for this state instruction."""
+        return "CD"
+
+    def get_arguments(self):
+        """Gets this state instruction's argument list."""
+        return [self.source_id, self.key, self.target_id]
+
+class DeleteEdgeInstruction(StateInstruction):
+    """An instruction that deletes an edge."""
+
+    def __init__(self, edge_id):
+        StateInstruction.__init__(self)
+        self.edge_id = edge_id
+
+    def simplify(self):
+        """Applies basic simplification to this instruction and its children."""
+        return DeleteEdgeInstruction(self.edge_id.simplify())
+
+    def has_result(self):
+        """Tells if this instruction computes a result."""
+        return False
+
+    def get_opcode(self):
+        """Gets the opcode for this state instruction."""
+        return "DE"
+
+    def get_arguments(self):
+        """Gets this state instruction's argument list."""
+        return [self.edge_id]
+
+def create_block(*statements):
+    """Creates a block-statement from the given list of statements."""
+    length = len(statements)
+    if length == 0:
+        return EmptyInstruction()
+    elif length == 1:
+        return statements[0]
+    else:
+        return CompoundInstruction(
+            statements[0],
+            create_block(*statements[1:]))
+
 if __name__ == "__main__":
     example_tree = SelectInstruction(
         LiteralInstruction(True),