Use the bytecode IR interpreter in adaptive-jit

kernel/modelverse_jit/jit.py

@@ -522,6 +522,7 @@ class ModelverseJit(object):
         del self.compilation_dependencies[body_id]
 
         if self.jit_success_log_function is not None:
+            assert self.jitted_entry_points[body_id] == function_name
             self.jit_success_log_function(
                 "JIT compilation successful: (function '%s' at %d)" % (function_name, body_id))
 
@@ -826,8 +827,11 @@ def favor_small_functions(body_bytecode):
 
 ADAPTIVE_JIT_LOOP_INSTRUCTION_MULTIPLIER = 4
 
-ADAPTIVE_FAST_JIT_TEMPERATURE_THRESHOLD = 200
-"""The threshold temperature at which fast-jit will be used."""
+ADAPTIVE_BASELINE_JIT_TEMPERATURE_THRESHOLD = 100
+"""The threshold temperature at which the adaptive JIT will use the baseline JIT."""
+
+ADAPTIVE_FAST_JIT_TEMPERATURE_THRESHOLD = 250
+"""The threshold temperature at which the adaptive JIT will use the fast JIT."""
 
 def favor_loops(body_bytecode):
     """Computes the initial temperature of a function. Code within a loop makes
@@ -857,6 +861,36 @@ class AdaptiveJitState(object):
         self.temperature_increment = temperature_increment
         self.can_rejit_name = can_rejit_name
 
+    def compile_interpreter(
+            self, jit, function_name, body_id, task_root):
+        """Compiles the given function as a function that controls the temperature counter
+           and calls the interpreter."""
+        def __increment_temperature(**kwargs):
+            if jit.jit_globals[self.can_rejit_name]:
+                temperature_counter_val = jit.jit_globals[self.temperature_counter_name]
+                temperature_counter_val += self.temperature_increment
+                jit.jit_globals[self.temperature_counter_name] = temperature_counter_val
+                if temperature_counter_val >= ADAPTIVE_BASELINE_JIT_TEMPERATURE_THRESHOLD:
+                    if temperature_counter_val >= ADAPTIVE_FAST_JIT_TEMPERATURE_THRESHOLD:
+                        yield [
+                            ("CALL_ARGS",
+                             [jit.jit_rejit,
+                              (task_root, body_id, function_name, compile_function_body_fast)])]
+                    else:
+                        yield [
+                            ("CALL_ARGS",
+                             [jit.jit_rejit,
+                              (task_root, body_id, function_name, self.compile_baseline)])]
+                    result, = yield [("CALL_KWARGS", [jit.jit_globals[function_name], kwargs])]
+                    raise primitive_functions.PrimitiveFinished((True, result))
+
+            raise primitive_functions.PrimitiveFinished((False, None))
+
+        yield [
+            ("TAIL_CALL_ARGS",
+             [compile_function_body_interpret,
+              (jit, function_name, body_id, task_root, __increment_temperature)])]
+
     def compile_baseline(
             self, jit, function_name, body_id, task_root):
         """Compiles the given function with the baseline JIT, and inserts logic that controls
@@ -932,11 +966,14 @@ def compile_function_body_adaptive(
     initial_temperature, temperature_increment = temperature_heuristic(body_bytecode)
     if jit.jit_success_log_function is not None:
         jit.jit_success_log_function(
-            'Initial temperature for %s: %d' % (function_name, initial_temperature))
+            "Initial temperature for '%s': %d" % (function_name, initial_temperature))
 
     if initial_temperature >= ADAPTIVE_FAST_JIT_TEMPERATURE_THRESHOLD:
         # Initial temperature exceeds the fast-jit threshold.
         # Compile this thing with fast-jit right away.
+        if jit.jit_success_log_function is not None:
+            jit.jit_success_log_function(
+                "Compiling '%s' with fast-jit." % function_name)
         yield [
             ("TAIL_CALL_ARGS",
              [compile_function_body_fast, (jit, function_name, body_id, task_root)])]
@@ -948,4 +985,21 @@ def compile_function_body_adaptive(
     jit.jit_globals[can_rejit_name] = True
 
     state = AdaptiveJitState(temperature_counter_name, temperature_increment, can_rejit_name)
-    yield [("TAIL_CALL_ARGS", [state.compile_baseline, (jit, function_name, body_id, task_root)])]
+
+    if initial_temperature >= ADAPTIVE_BASELINE_JIT_TEMPERATURE_THRESHOLD:
+        # Initial temperature exceeds the baseline JIT threshold.
+        # Compile this thing with baseline JIT right away.
+        if jit.jit_success_log_function is not None:
+            jit.jit_success_log_function(
+                "Compiling '%s' with baseline-jit." % function_name)
+        yield [
+            ("TAIL_CALL_ARGS",
+             [state.compile_baseline, (jit, function_name, body_id, task_root)])]
+    else:
+        # Looks like we'll use the interpreter initially.
+        if jit.jit_success_log_function is not None:
+            jit.jit_success_log_function(
+                "Compiling '%s' with bytecode-interpreter." % function_name)
+        yield [
+            ("TAIL_CALL_ARGS",
+             [state.compile_interpreter, (jit, function_name, body_id, task_root)])]