Various improvements to JIT

kernel/modelverse_jit/jit.py

@@ -1,6 +1,7 @@
 import modelverse_kernel.primitives as primitive_functions
 import modelverse_jit.tree_ir as tree_ir
 import modelverse_jit.runtime as jit_runtime
+import keyword
 
 KWARGS_PARAMETER_NAME = "kwargs"
 """The name of the kwargs parameter in jitted functions."""
@@ -153,7 +154,9 @@ class ModelverseJit(object):
     def generate_function_name(self, suggested_name=None):
         """Generates a new function name or picks the suggested name if it is still
            available."""
-        if suggested_name is not None and suggested_name not in self.jit_globals:
+        if suggested_name is not None \
+            and suggested_name not in self.jit_globals \
+            and not keyword.iskeyword(suggested_name):
             self.jit_count += 1
             return suggested_name
         else:
@@ -328,6 +331,7 @@ class AnalysisState(object):
         self.user_root = user_root
         self.jit = jit
         self.local_mapping = local_mapping
+        self.function_name = jit.jitted_entry_points[body_id]
 
     def get_local_name(self, local_id):
         """Gets the name for a local with the given id."""
@@ -400,7 +404,8 @@ 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)
+                outer_result = tree_ir.with_debug_info_trace(
+                    outer_e.result, debug_info, self.function_name)
                 if next_instr is None:
                     raise primitive_functions.PrimitiveFinished(outer_result)
                 else:

kernel/modelverse_jit/tree_ir.py

@@ -1152,15 +1152,19 @@ def create_block(*statements):
             statements[0],
             create_block(*statements[1:]))
 
-def with_debug_info_trace(instruction, debug_info):
+def with_debug_info_trace(instruction, debug_info, function_name):
     """Prepends the given instruction with a tracing instruction that prints
-       the given debug information."""
-    if debug_info is None:
+       the given debug information and function name."""
+    if debug_info is None and function_name is None:
         return instruction
     else:
+        if debug_info is None:
+            debug_info = 'unknown location'
+        if function_name is None:
+            function_name = 'unknown function'
         return create_block(
             PrintInstruction(
-                LiteralInstruction('TRACE: %s(JIT)' % debug_info)),
+                LiteralInstruction('TRACE: %s(%s, JIT)' % (debug_info, function_name))),
             instruction)
 
 def map_and_simplify(function, instruction):