Add support for JIT tracing

kernel/modelverse_jit/jit.py

@@ -57,6 +57,7 @@ class ModelverseJit(object):
         self.compilation_dependencies = {}
         self.jit_enabled = True
         self.direct_calls_allowed = True
+        self.tracing_enabled = False
 
     def set_jit_enabled(self, is_enabled=True):
         """Enables or disables the JIT."""
@@ -66,6 +67,10 @@ class ModelverseJit(object):
         """Allows or disallows direct calls from jitted to jitted code."""
         self.direct_calls_allowed = is_allowed
 
+    def enable_tracing(self, is_enabled=True):
+        """Enables or disables tracing for jitted code."""
+        self.tracing_enabled = is_enabled
+
     def mark_entry_point(self, body_id):
         """Marks the node with the given identifier as a function entry point."""
         if body_id not in self.no_jit_entry_points and body_id not in self.jitted_entry_points:
@@ -322,6 +327,15 @@ class AnalysisState(object):
         instruction_val, = yield [("RV", [instruction_id])]
         instruction_val = instruction_val["value"]
         if instruction_val in self.instruction_analyzers:
+            # If tracing is enabled, then this would be an appropriate time to
+            # retrieve the debug information.
+            if self.jit.tracing_enabled:
+                debug_info, = yield [("RD", [instruction_id, "__debug"])]
+                if debug_info is not None:
+                    debug_info, = yield [("RV", [debug_info])]
+            else:
+                debug_info = None
+
             try:
                 gen = self.instruction_analyzers[instruction_val](self, instruction_id)
                 inp = None
@@ -334,8 +348,9 @@ class AnalysisState(object):
             except primitive_functions.PrimitiveFinished as outer_e:
                 # Check if the instruction has a 'next' instruction.
                 next_instr, = yield [("RD", [instruction_id, "next"])]
+                outer_result = tree_ir.with_debug_info_trace(outer_e.result, debug_info)
                 if next_instr is None:
-                    raise outer_e
+                    raise primitive_functions.PrimitiveFinished(outer_result)
                 else:
                     gen = self.analyze(next_instr)
                     try:
@@ -345,7 +360,7 @@ class AnalysisState(object):
                     except primitive_functions.PrimitiveFinished as inner_e:
                         raise primitive_functions.PrimitiveFinished(
                             tree_ir.CompoundInstruction(
-                                outer_e.result,
+                                outer_result,
                                 inner_e.result))
         else:
             raise JitCompilationFailedException(

kernel/modelverse_jit/tree_ir.py

@@ -284,7 +284,7 @@ class CallInstruction(Instruction):
                 arg.generate_python_def(code_generator)
 
         code_generator.append_line(
-            '%s = %s(%s) ' % (
+            '%s = %s(%s)' % (
                 code_generator.get_result_name(self),
                 self.target.generate_python_use(code_generator),
                 ', '.join([arg.generate_python_use(code_generator) for arg in self.argument_list])))
@@ -323,23 +323,49 @@ class JitCallInstruction(Instruction):
         arg_list.append(
             '**%s' % self.kwarg.generate_python_use(code_generator))
 
+        own_name = code_generator.get_result_name(self)
         code_generator.append_line('try:')
         code_generator.increase_indentation()
         code_generator.append_line(
-            'gen = %s(%s) ' % (
+            '%s_gen = %s(%s)' % (
+                own_name,
                 self.target.generate_python_use(code_generator),
                 ', '.join(arg_list)))
-        code_generator.append_line('inp = None')
+        code_generator.append_line('%s_inp = None' % own_name)
         code_generator.append_line('while 1:')
         code_generator.increase_indentation()
-        code_generator.append_line('inp = yield gen.send(inp)')
+        code_generator.append_line(
+            '%s_inp = yield %s_gen.send(%s_inp)' % (own_name, own_name, own_name))
         code_generator.decrease_indentation()
         code_generator.decrease_indentation()
-        code_generator.append_line('except PrimitiveFinished as ex:')
+        code_generator.append_line('except PrimitiveFinished as %s_ex:' % own_name)
         code_generator.increase_indentation()
-        code_generator.append_line('%s = ex.result' % code_generator.get_result_name(self))
+        code_generator.append_line('%s = %s_ex.result' % (own_name, own_name))
         code_generator.decrease_indentation()
 
+class PrintInstruction(Instruction):
+    """An instruction that prints a value."""
+    def __init__(self, argument):
+        Instruction.__init__(self)
+        self.argument = argument
+
+    def has_result(self):
+        """Tells if this instruction has a result."""
+        return False
+
+    def simplify(self):
+        """Applies basic simplification to this instruction and its children."""
+        return PrintInstruction(self.argument.simplify())
+
+    def generate_python_def(self, code_generator):
+        """Generates Python code for this instruction."""
+        if self.argument.has_definition():
+            self.argument.generate_python_def(code_generator)
+
+        code_generator.append_line(
+            'print(%s)' % (
+                self.argument.generate_python_use(code_generator)))
+
 class BinaryInstruction(Instruction):
     """An instruction that performs a binary operation."""
     def __init__(self, lhs, operator, rhs):
@@ -958,6 +984,17 @@ def create_block(*statements):
             statements[0],
             create_block(*statements[1:]))
 
+def with_debug_info_trace(instruction, debug_info):
+    """Prepends the given instruction with a tracing instruction that prints
+       the given debug information."""
+    if debug_info is None:
+        return instruction
+    else:
+        return create_block(
+            PrintInstruction(
+                LiteralInstruction('TRACE: %s(JIT)' % debug_info)),
+            instruction)
+
 if __name__ == "__main__":
     example_tree = SelectInstruction(
         LiteralInstruction(True),

kernel/modelverse_kernel/main.py

@@ -38,6 +38,11 @@ class ModelverseKernel(object):
         #
         # To disable direct calls in the JIT, uncomment the line below:
         #     self.jit.allow_direct_calls(False)
+        #
+        # To enable tracing in the JIT (for debugging purposes), uncomment
+        # the line below:
+        #
+        #     self.jit.enable_tracing()
         self.debug_info = defaultdict(list)
 
     def execute_yields(self, username, operation, params, reply):