Support compiled function lookup in the JIT

kernel/modelverse_jit/jit.py

@@ -11,7 +11,7 @@ class JitCompilationFailedException(Exception):
 class ModelverseJit(object):
     """A high-level interface to the modelverse JIT compiler."""
 
-    def __init__(self, max_instructions=None):
+    def __init__(self, max_instructions=None, compiled_function_lookup=None):
         self.todo_entry_points = set()
         self.no_jit_entry_points = set()
         self.jitted_entry_points = {}
@@ -21,6 +21,7 @@ class ModelverseJit(object):
         }
         self.jit_count = 0
         self.max_instructions = 30 if max_instructions is None else max_instructions
+        self.compiled_function_lookup = compiled_function_lookup
         self.jit_enabled = True
 
     def set_jit_enabled(self, is_enabled=True):
@@ -78,6 +79,16 @@ class ModelverseJit(object):
         if body_id in self.todo_entry_points:
             self.todo_entry_points.remove(body_id)
 
+    def lookup_compiled_function(self, name):
+        """Looks up a compiled function by name. Returns a matching function,
+           or None if no function was found."""
+        if name in self.jit_globals:
+            return self.jit_globals[name]
+        elif self.compiled_function_lookup is not None:
+            return self.compiled_function_lookup(name)
+        else:
+            return None
+
     def jit_parameters(self, body_id):
         """Acquires the parameter list for the given body id node."""
         if body_id not in self.jitted_parameters:
@@ -654,7 +665,6 @@ class AnalysisState(object):
 
     def analyze_direct_call(self, callee_id, callee_name, first_parameter_id):
         """Tries to analyze a direct 'call' instruction."""
-
         self.register_function_var(callee_id)
 
         body_id, = yield [("RD", [callee_id, "body"])]
@@ -667,14 +677,18 @@ class AnalysisState(object):
         except primitive_functions.PrimitiveFinished as ex:
             _, parameter_names = ex.result
 
-        # Compile the callee.
-        try:
-            gen = self.jit.jit_compile(self.user_root, body_id)
-            inp = None
-            while True:
-                inp = yield gen.send(inp)
-        except primitive_functions.PrimitiveFinished as ex:
-            pass
+        compiled_func = self.jit.lookup_compiled_function(callee_name)
+        if compiled_func is None:
+            # Compile the callee.
+            try:
+                gen = self.jit.jit_compile(self.user_root, body_id)
+                inp = None
+                while True:
+                    inp = yield gen.send(inp)
+            except primitive_functions.PrimitiveFinished as ex:
+                pass
+        else:
+            self.jit.register_compiled(body_id, compiled_func, callee_name)
 
         # Get the callee's name.
         compiled_func_name = self.jit.get_compiled_name(body_id)