Postpone CFG JIT intrinsic expansion in the pass pipeline

kernel/modelverse_jit/cfg_data_structures.py

@@ -50,13 +50,21 @@ def get_call_def_rewriter(definition, def_rewrite_rules):
 
 def get_call_use_rewriter(use_definition, def_definition, use_rewrite_rules):
     """Gets an appropriate rewrite rule from the given dictionary of call rewrite rules."""
-    if cfg_ir.is_call(use_definition, calling_convention=cfg_ir.JIT_CALLING_CONVENTION):
-        call = cfg_ir.get_def_value(use_definition)
-        for arg_name, arg_def in call.argument_list:
-            if arg_def == def_definition:
-                key = (call.target_name, arg_name)
-                if key in use_rewrite_rules:
-                    return use_rewrite_rules[key]
+    if not cfg_ir.is_call(use_definition):
+        return None
+
+    call = cfg_ir.get_def_value(use_definition)
+    if call.calling_convention not in (
+            cfg_ir.JIT_CALLING_CONVENTION,
+            cfg_ir.JIT_NO_GC_CALLING_CONVENTION,
+            cfg_ir.JIT_CFG_INTRINSIC_CALLING_CONVENTION):
+        return None
+
+    for arg_name, arg_def in call.argument_list:
+        if arg_def == def_definition:
+            key = (call.target_name, arg_name)
+            if key in use_rewrite_rules:
+                return use_rewrite_rules[key]
 
     return None
 

kernel/modelverse_jit/cfg_ir.py

@@ -486,6 +486,9 @@ SELF_POSITIONAL_CALLING_CONVENTION = 'self-positional'
 JIT_CALLING_CONVENTION = 'jit'
 """The calling convention for jitted functions."""
 
+JIT_CFG_INTRINSIC_CALLING_CONVENTION = 'jit-cfg-intrinsic'
+"""The calling convention for CFG JIT intrinsics."""
+
 JIT_NO_GC_CALLING_CONVENTION = 'jit-no-gc'
 """The calling convention for jitted functions that promise not to initiate a GC cycle."""
 
@@ -510,10 +513,10 @@ NOP_MACRO_NAME = 'nop'
 READ_DICT_KEYS_MACRO_NAME = 'read_dict_keys'
 """The name of the macro that reads all keys from a dictionary."""
 
-READ_DICT_VALUE_MACRO_NAME = "read_dict_value"
+READ_DICT_VALUE_MACRO_NAME = 'read_dict_value'
 """The name of the macro that reads the value for a given key from a dictionary."""
 
-READ_DICT_NODE_MACRO_NAME = "read_dict_node"
+READ_DICT_NODE_MACRO_NAME = 'read_dict_node'
 """The name of the macro that reads the node for a given key from a dictionary."""
 
 READ_OUTGOING_EDGES_MACRO_NAME = 'read_outgoing_edges'

kernel/modelverse_jit/cfg_optimization.py

@@ -237,7 +237,10 @@ def try_redefine_as_direct_call(definition, jit, called_globals):
             # Try to resolve the callee as an intrinsic.
             intrinsic = jit.get_cfg_intrinsic(resolved_var_name)
             if intrinsic is not None:
-                apply_cfg_intrinsic(intrinsic, definition, call.argument_list)
+                definition.redefine(
+                    cfg_ir.DirectFunctionCall(
+                        resolved_var_name, call.argument_list,
+                        cfg_ir.JIT_CFG_INTRINSIC_CALLING_CONVENTION))
             else:
                 # Otherwise, build a thunk.
                 thunk_name = jit.jit_thunk_global(resolved_var_name)
@@ -287,8 +290,7 @@ def optimize_calls(entry_point, jit):
        given entry point."""
     called_globals = set()
     global_exists_defs = defaultdict(list)
-    all_blocks = list(cfg_ir.get_all_blocks(entry_point))
-    for block in all_blocks:
+    for block in cfg_ir.get_all_blocks(entry_point):
         for definition in block.definitions:
             checked_global = get_checked_global(definition)
             if checked_global is not None:
@@ -300,6 +302,17 @@ def optimize_calls(entry_point, jit):
         for exists_def in global_exists_defs[resolve_global]:
             exists_def.redefine(cfg_ir.Literal(False))
 
+def expand_cfg_intrinsics(entry_point, jit):
+    """Expands CFG JIT intrinsics in the control-flow graph defined by the given entry point."""
+    for block in cfg_ir.get_all_blocks(entry_point):
+        for definition in block.definitions:
+            def_value = definition.value
+            if (isinstance(def_value, cfg_ir.DirectFunctionCall)
+                    and def_value.calling_convention ==
+                    cfg_ir.JIT_CFG_INTRINSIC_CALLING_CONVENTION):
+                intrinsic = jit.get_cfg_intrinsic(def_value.target_name)
+                apply_cfg_intrinsic(intrinsic, definition, def_value.argument_list)
+
 def simplify_values(entry_point):
     """Simplifies values in the control-flow graph defined by the given entry point."""
     for block in cfg_ir.get_all_blocks(entry_point):
@@ -514,6 +527,7 @@ def optimize(entry_point, jit):
     if jit.direct_calls_allowed:
         optimize_calls(entry_point, jit)
     cfg_data_structures.optimize_data_structures(entry_point)
+    expand_cfg_intrinsics(entry_point, jit)
     yield [("CALL_ARGS", [inline_constants, (entry_point,)])]
     optimize_reads(entry_point)
     simplify_values(entry_point)