Fold constant reads in the JIT

kernel/modelverse_jit/jit.py

@@ -34,6 +34,50 @@ def apply_intrinsic(intrinsic_function, named_args):
             tree_ir.create_block(*store_instructions),
             intrinsic_function(**arg_value_dict))
 
+def map_and_simplify_generator(function, instruction):
+    """Applies the given mapping function to every instruction in the tree
+       that has the given instruction as root, and simplifies it on-the-fly.
+
+       This is at least as powerful as first mapping and then simplifying, as
+       maps and simplifications are interspersed.
+
+       This function assumes that function creates a generator that returns by
+       raising a primitive_functions.PrimitiveFinished."""
+
+    # First handle the children by mapping on them and then simplifying them.
+    new_children = []
+    for inst in instruction.get_children():
+        try:
+            gen = map_and_simplify_generator(function, inst)
+            inp = None
+            while True:
+                inp = yield gen.send(inp)
+        except primitive_functions.PrimitiveFinished as ex:
+            new_children.append(ex.result)
+
+    # Then apply the function to the top-level node.
+    try:
+        gen = function(instruction.create(new_children))
+        inp = None
+        while True:
+            inp = yield gen.send(inp)
+    except primitive_functions.PrimitiveFinished as ex:
+        # Finally, simplify the transformed top-level node.
+        raise primitive_functions.PrimitiveFinished(ex.result.simplify_node())
+
+def expand_constant_read(instruction):
+    """Tries to replace a read of a constant node by a literal."""
+    if isinstance(instruction, tree_ir.ReadValueInstruction) and \
+        isinstance(instruction.node_id, tree_ir.LiteralInstruction):
+        val, = yield [("RV", [instruction.node_id.literal])]
+        raise primitive_functions.PrimitiveFinished(tree_ir.LiteralInstruction(val))
+    else:
+        raise primitive_functions.PrimitiveFinished(instruction)
+
+def optimize_tree_ir(instruction):
+    """Optimizes an IR tree."""
+    return map_and_simplify_generator(expand_constant_read, instruction)
+
 class JitCompilationFailedException(Exception):
     """A type of exception that is raised when the jit fails to compile a function."""
     pass
@@ -252,11 +296,19 @@ class ModelverseJit(object):
         constructed_body = tree_ir.create_block(
             *(prologue_statements + [constructed_body]))
 
+        try:
+            gen = optimize_tree_ir(constructed_body)
+            inp = None
+            while True:
+                inp = yield gen.send(inp)
+        except primitive_functions.PrimitiveFinished as ex:
+            constructed_body = ex.result
+
         # Wrap the IR in a function definition, give it a unique name.
         constructed_function = tree_ir.DefineFunctionInstruction(
             function_name,
             parameter_list + ['**' + KWARGS_PARAMETER_NAME],
-            constructed_body.simplify())
+            constructed_body)
         # Convert the function definition to Python code, and compile it.
         exec(str(constructed_function), self.jit_globals)
         # Extract the compiled function from the JIT global state.
@@ -370,8 +422,8 @@ class AnalysisState(object):
         """Tries to compile a list of IR trees from the given list of instruction ids."""
         results = []
         for inst in instruction_ids:
-            gen = self.analyze(inst)
             try:
+                gen = self.analyze(inst)
                 inp = None
                 while True:
                     inp = yield gen.send(inp)
@@ -388,8 +440,8 @@ class AnalysisState(object):
                 tree_ir.ReturnInstruction(
                     tree_ir.EmptyInstruction()))
         else:
-            gen = self.analyze(retval_id)
             try:
+                gen = self.analyze(retval_id)
                 inp = None
                 while True:
                     inp = yield gen.send(inp)

kernel/modelverse_jit/tree_ir.py

@@ -1129,6 +1129,18 @@ def with_debug_info_trace(instruction, debug_info):
                 LiteralInstruction('TRACE: %s(JIT)' % debug_info)),
             instruction)
 
+def map_and_simplify(function, instruction):
+    """Applies the given mapping function to every instruction in the tree
+       that has the given instruction as root, and simplifies it on-the-fly.
+
+       This is at least as powerful as first mapping and then simplifying, as
+       maps and simplifications are interspersed."""
+    # Let's just agree to disagree on map vs list comprehensions, pylint.
+    # pylint: disable=I0011,W0141
+    return function(
+        instruction.create(
+            map(map_and_simplify, instruction.get_children()))).simplify_node()
+
 if __name__ == "__main__":
     example_tree = SelectInstruction(
         LiteralInstruction(True),