Tweak temperature computation in the adaptive JIT

kernel/modelverse_jit/jit.py

@@ -718,13 +718,18 @@ class ModelverseJit(object):
                 tree_ir.LiteralInstruction(jit_runtime.FUNCTION_BODY_KEY)),
             global_name)
 
-def compile_function_body_interpret(jit, function_name, body_id, task_root):
+def compile_function_body_interpret(jit, function_name, body_id, task_root, header=None):
     """Create a function that invokes the interpreter on the given function."""
     (parameter_ids, parameter_list, _), = yield [
         ("CALL_ARGS", [jit.jit_signature, (body_id,)])]
     param_dict = dict(zip(parameter_ids, parameter_list))
     body_bytecode, = yield [("CALL_ARGS", [jit.jit_parse_bytecode, (body_id,)])]
     def __interpret_function(**kwargs):
+        if header is not None:
+            (done, result), = yield [("CALL_KWARGS", [header, kwargs])]
+            if done:
+                raise primitive_functions.PrimitiveFinished(result)
+
         local_args = {}
         inner_kwargs = dict(kwargs)
         for param_id, name in param_dict.items():
@@ -806,13 +811,7 @@ def favor_large_functions(body_bytecode):
     # in a situation where said function runs for a long time before we
     # realize that we really should have jitted it. And that's exactly what
     # this heuristic tries to avoid.
-    return (
-        len(body_bytecode.get_reachable()),
-        lambda old_value:
-        tree_ir.BinaryInstruction(
-            old_value,
-            '+',
-            tree_ir.LiteralInstruction(1)))
+    return len(body_bytecode.get_reachable()), 1
 
 def favor_small_functions(body_bytecode):
     """Computes the initial temperature of a function based on the size of
@@ -823,13 +822,7 @@ def favor_small_functions(body_bytecode):
     # of fast-jit's algorithms. So it might be cheaper to fast-jit small
     # functions and get a performance boost from that than to fast-jit large
     # functions.
-    return (
-        ADAPTIVE_FAST_JIT_TEMPERATURE_THRESHOLD - len(body_bytecode.get_reachable()),
-        lambda old_value:
-        tree_ir.BinaryInstruction(
-            old_value,
-            '+',
-            tree_ir.LiteralInstruction(1)))
+    return ADAPTIVE_FAST_JIT_TEMPERATURE_THRESHOLD - len(body_bytecode.get_reachable()), 1
 
 ADAPTIVE_JIT_LOOP_INSTRUCTION_MULTIPLIER = 4
 
@@ -853,13 +846,7 @@ def favor_loops(body_bytecode):
                     x, (bytecode_ir.BreakInstruction, bytecode_ir.ContinueInstruction)))
             temperature += ADAPTIVE_JIT_LOOP_INSTRUCTION_MULTIPLIER * len(loop_body_instructions)
 
-    return (
-        temperature,
-        lambda old_value:
-        tree_ir.BinaryInstruction(
-            old_value,
-            '+',
-            tree_ir.LiteralInstruction(1)))
+    return temperature, 1
 
 def compile_function_body_adaptive(
         jit, function_name, body_id, task_root,
@@ -872,7 +859,7 @@ def compile_function_body_adaptive(
     # and gets incremented every time the function is executed.
 
     body_bytecode, = yield [("CALL_ARGS", [jit.jit_parse_bytecode, (body_id,)])]
-    initial_temperature, increment_temperature = temperature_heuristic(body_bytecode)
+    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))
@@ -897,7 +884,7 @@ def compile_function_body_adaptive(
     #
     # if can_rejit:
     #     global temperature_counter
-    #     temperature_counter = increment_temperature(temperature_counter)
+    #     temperature_counter = temperature_counter + temperature_increment
     #     if temperature_counter >= ADAPTIVE_FAST_JIT_TEMPERATURE_THRESHOLD:
     #         yield [("CALL_KWARGS", [jit_runtime.JIT_REJIT_FUNCTION_NAME, {...}])]
     #         yield [("TAIL_CALL_KWARGS", [function_name, {...}])]
@@ -909,8 +896,10 @@ def compile_function_body_adaptive(
             tree_ir.IgnoreInstruction(
                 tree_ir.StoreGlobalInstruction(
                     temperature_counter_name,
-                    increment_temperature(
-                        tree_ir.LoadGlobalInstruction(temperature_counter_name)))),
+                    tree_ir.BinaryInstruction(
+                        tree_ir.LoadGlobalInstruction(temperature_counter_name),
+                        '+',
+                        tree_ir.LiteralInstruction(temperature_increment)))),
             tree_ir.SelectInstruction(
                 tree_ir.BinaryInstruction(
                     tree_ir.LoadGlobalInstruction(temperature_counter_name),