Define CFG JIT intrinsics

kernel/modelverse_jit/jit.py

@@ -5,6 +5,7 @@ import modelverse_jit.bytecode_to_tree as bytecode_to_tree
 import modelverse_jit.bytecode_to_cfg as bytecode_to_cfg
 import modelverse_jit.cfg_optimization as cfg_optimization
 import modelverse_jit.cfg_to_tree as cfg_to_tree
+import modelverse_jit.cfg_ir as cfg_ir
 import modelverse_jit.tree_ir as tree_ir
 import modelverse_jit.runtime as jit_runtime
 
@@ -118,6 +119,8 @@ class ModelverseJit(object):
         self.compiled_function_lookup = compiled_function_lookup
         # jit_intrinsics is a function name -> intrinsic map.
         self.jit_intrinsics = {}
+        # cfg_jit_intrinsics is a function name -> intrinsic map.
+        self.cfg_jit_intrinsics = {}
         self.compilation_dependencies = {}
         self.jit_enabled = True
         self.direct_calls_allowed = True
@@ -317,18 +320,39 @@ class ModelverseJit(object):
         else:
             return None
 
-    def register_intrinsic(self, name, intrinsic_function):
+    def get_cfg_intrinsic(self, name):
+        """Tries to find an intrinsic version of the function with the
+           given name that is specialized for CFGs."""
+        if name in self.cfg_jit_intrinsics:
+            return self.cfg_jit_intrinsics[name]
+        else:
+            return None
+
+    def register_intrinsic(self, name, intrinsic_function, cfg_intrinsic_function=None):
         """Registers the given intrisic with the JIT. This will make the JIT replace calls to
            the function with the given entry point by an application of the specified function."""
         self.jit_intrinsics[name] = intrinsic_function
+        if cfg_intrinsic_function is not None:
+            self.cfg_jit_intrinsics[name] = cfg_intrinsic_function
 
     def register_binary_intrinsic(self, name, operator):
         """Registers an intrinsic with the JIT that represents the given binary operation."""
-        self.register_intrinsic(name, lambda a, b: tree_ir.CreateNodeWithValueInstruction(
-            tree_ir.BinaryInstruction(
-                tree_ir.ReadValueInstruction(a),
-                operator,
-                tree_ir.ReadValueInstruction(b))))
+        self.register_intrinsic(
+            name,
+            lambda a, b:
+            tree_ir.CreateNodeWithValueInstruction(
+                tree_ir.BinaryInstruction(
+                    tree_ir.ReadValueInstruction(a),
+                    operator,
+                    tree_ir.ReadValueInstruction(b))),
+            lambda original_def, a, b:
+            original_def.redefine(
+                cfg_ir.CreateNode(
+                    original_def.insert_before(
+                        cfg_ir.Binary(
+                            original_def.insert_before(cfg_ir.Read(a)),
+                            operator,
+                            original_def.insert_before(cfg_ir.Read(b)))))))
 
     def register_unary_intrinsic(self, name, operator):
         """Registers an intrinsic with the JIT that represents the given unary operation."""