Embed logging hooks in the JIT

kernel/modelverse_jit/jit.py

@@ -82,6 +82,10 @@ def optimize_tree_ir(instruction):
     """Optimizes an IR tree."""
     return map_and_simplify_generator(expand_constant_read, instruction)
 
+def print_value(val):
+    """A thin wrapper around 'print'."""
+    print(val)
+
 class ModelverseJit(object):
     """A high-level interface to the modelverse JIT compiler."""
     def __init__(self, max_instructions=None, compiled_function_lookup=None):
@@ -105,6 +109,8 @@ class ModelverseJit(object):
         self.tracing_enabled = False
         self.input_function_enabled = False
         self.nop_insertion_enabled = True
+        self.jit_success_log_function = None
+        self.jit_code_log_function = None
 
     def set_jit_enabled(self, is_enabled=True):
         """Enables or disables the JIT."""
@@ -129,6 +135,16 @@ class ModelverseJit(object):
            Modelverse server an opportunity to interrupt the currently running code."""
         self.nop_insertion_enabled = is_enabled
 
+    def set_jit_success_log(self, log_function=print_value):
+        """Configures this JIT instance with a function that prints output to a log.
+           Success and failure messages for specific functions are then sent to said log."""
+        self.jit_success_log_function = log_function
+
+    def set_jit_code_log(self, log_function=print_value):
+        """Configures this JIT instance with a function that prints output to a log.
+           Function definitions of jitted functions are then sent to said log."""
+        self.jit_code_log_function = log_function
+
     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:
@@ -306,8 +322,11 @@ class ModelverseJit(object):
                 if dep in self.jitted_entry_points:
                     del self.jitted_entry_points[dep]
 
-            raise JitCompilationFailedException(
-                "%s (function '%s' at %d)" % (exception.message, function_name, body_id))
+            failure_message = "%s (function '%s' at %d)" % (
+                exception.message, function_name, body_id)
+            if self.jit_success_log_function is not None:
+                self.jit_success_log_function('JIT compilation failed: %s' % failure_message)
+            raise JitCompilationFailedException(failure_message)
 
         # Try to analyze the function's body.
         yield [("TRY", [])]
@@ -360,7 +379,13 @@ class ModelverseJit(object):
         # Extract the compiled function from the JIT global state.
         compiled_function = self.jit_globals[function_name]
 
-        # print(constructed_function)
+        if self.jit_success_log_function is not None:
+            self.jit_success_log_function(
+                "JIT compilation successful: (function '%s' at %d)" % (function_name, body_id))
+
+        if self.jit_code_log_function is not None:
+            self.jit_code_log_function(constructed_function)
+
         raise primitive_functions.PrimitiveFinished(compiled_function)
 
 class AnalysisState(object):

kernel/modelverse_kernel/main.py

@@ -52,15 +52,33 @@ class ModelverseKernel(object):
         #
         #     self.jit.allow_direct_calls(False)
         #
+        # To make the JIT compile 'input' instructions as calls to
+        # modelverse_jit.runtime.get_input, uncomment the line below:
+        #
+        #     self.jit.use_input_function()
+        #
         # To enable tracing in the JIT (for debugging purposes), uncomment
         # the line below:
         #
         #     self.jit.enable_tracing()
         #
-        # To make the JIT compile 'input' instructions as calls to
-        # modelverse_jit.runtime.get_input, uncomment the line below:
+        # To make the JIT print JIT successes and errors to the command-line,
+        # uncomment the line below:
         #
-        #     self.jit.use_input_function()
+        #     self.jit.set_jit_success_log()
+        #
+        # If you want, you can use a custom logging function:
+        #
+        #     self.jit.set_jit_success_log(logging_function)
+        #
+        # To make the JIT print jitted code to the command-line, uncomment the
+        # line below:
+        #
+        #     self.jit.set_jit_code_log()
+        #
+        # If you want, you can use a custom logging function:
+        #
+        #     self.jit.set_jit_code_log(logging_function)
         #
 
         self.debug_info = defaultdict(list)
@@ -122,7 +140,6 @@ class ModelverseKernel(object):
 
         def handle_jit_failed(exception):
             # Try again, but this time without the JIT.
-            # print(exception.message)
             gen = self.get_inst_phase_generator(inst_v, self.phase_v, user_root)
             yield [("TAIL_CALL", [gen])]