Allow inlinable LLVM function definitions.

Author: braxtonmckee

typed_python/compiler/native_compiler/native_ast.py

@@ -208,7 +208,8 @@ def const_str(c):
     external=bool,
     varargs=bool,
     intrinsic=bool,
-    can_throw=bool
+    can_throw=bool,
+    inlineLlvmDefinition=OneOf(None, str)
 )
 
 

typed_python/compiler/native_compiler/native_ast_to_llvm.py

@@ -24,6 +24,7 @@
 
 import typed_python.compiler.native_compiler.native_ast as native_ast
 import llvmlite.ir
+import re
 import os
 
 
@@ -293,6 +294,7 @@ def add_functions(self, names_to_definitions):
 
         globalDefinitions = {}
         globalDefinitionsLlvmValues = {}
+        extraDefinitions = {}
 
         while names_to_definitions:
             for name in sorted(names_to_definitions):
@@ -322,7 +324,8 @@ def add_functions(self, names_to_definitions):
                             builder,
                             arg_assignments,
                             definition.output_type,
-                            external_function_references
+                            external_function_references,
+                            extraDefinitions
                         )
 
                         func_converter.setup()
@@ -375,8 +378,21 @@ def add_functions(self, names_to_definitions):
             if name not in self._function_definitions:
                 raise Exception(f"Somehow we depend on {name} but have no definition for it")
 
+        moduleText = str(module)
+
+        # substitute in direct LLVM IR function definitions for the 'extraDefinitins'
+        # llvmlite produces 'declare' stubs for these functions which we need to remove
+        # so that we can place a direct declaration. This is obviously very fragile but
+        # gives us an ability to work around llvmlite not having total coverage for all
+        # llvm features we might want (such as vectors).
+        for fName, fDef in extraDefinitions.items():
+            moduleText = moduleText + "\n\n" + fDef
+
+            pattern = '\ndeclare .*@"' + fName + '"\\([^)]*\\)'
+            moduleText = re.sub(pattern, '', moduleText)
+
         return ModuleDefinition(
-            str(module),
+            moduleText,
             functionTypes,
             globalDefinitions,
             usedExternalFunctions,

typed_python/compiler/native_compiler/native_ast_to_llvm_function_converter.py

@@ -567,7 +567,8 @@ def __init__(self,
                  builder,
                  arg_assignments,
                  output_type,
-                 external_function_references
+                 external_function_references,
+                 extraDefinitions
                  ):
         self.function = function
 
@@ -583,6 +584,7 @@ def __init__(self,
         self.external_function_references = external_function_references
         self.tags_initialized = {}
         self.stack_slots = {}
+        self.extraDefinitions = extraDefinitions
 
     def tags_as(self, new_tags):
         class scoper():
@@ -694,6 +696,9 @@ def generate_exception_and_store_value(self, llvm_pointer_val):
 
     def namedCallTargetToLLVM(self, target):
         if target.external:
+            if target.inlineLlvmDefinition:
+                self.extraDefinitions[target.name] = target.inlineLlvmDefinition
+
             if target.name not in self.external_function_references:
                 func_type = llvmlite.ir.FunctionType(
                     type_to_llvm_type(target.output_type),

typed_python/compiler/tests/compilable_builtin_test.py

@@ -0,0 +1,36 @@
+from typed_python import Entrypoint
+from typed_python.compiler.type_wrappers.compilable_builtin import CompilableBuiltin
+from typed_python.compiler.type_wrappers.runtime_functions import externalCallTarget, Float64
+
+
+inlineLlvmFunc = externalCallTarget("inlineLlvmFunc", Float64, Float64, inlineLlvmDefinition="""
+    define external double @"inlineLlvmFunc"(double %".1") {
+    entry:
+      %.4 = fadd double %.1, 1.000000e+00
+      ret double %.4
+    }
+""")
+
+
+class InlineLlvmFunc(CompilableBuiltin):
+    def __eq__(self, other):
+        return isinstance(other, inlineLlvmFunc)
+
+    def __hash__(self):
+        return hash("inlineLlvmFunc")
+
+    def convert_call(self, context, instance, args, kwargs):
+        return context.pushPod(
+            float,
+            inlineLlvmFunc.call(
+                args[0],
+            )
+        )
+
+
+def test_inline_llvm():
+    @Entrypoint
+    def f(x):
+        return InlineLlvmFunc()(x)
+
+    assert f(2.5) == 3.5

typed_python/compiler/type_wrappers/runtime_functions.py

@@ -31,7 +31,7 @@
 Void = native_ast.Void
 
 
-def externalCallTarget(name, output, *inputs, varargs=False, intrinsic=False, canThrow=False):
+def externalCallTarget(name, output, *inputs, varargs=False, intrinsic=False, canThrow=False, inlineLlvmDefinition=None):
     """Create an object that lets us call C functions.
 
     Note that 'canThrow' really indicates whether we should use llvm 'invoke' instead
@@ -46,7 +46,8 @@ def externalCallTarget(name, output, *inputs, varargs=False, intrinsic=False, ca
             external=True,
             varargs=varargs,
             intrinsic=intrinsic,
-            can_throw=canThrow
+            can_throw=canThrow,
+            inlineLlvmDefinition=inlineLlvmDefinition
         )
     )