NativeCompiler computes the link-dependencies of the functions provided to it.

Author: braxtonmckee

typed_python/compiler/native_compiler/binary_shared_object.py

@@ -36,7 +36,7 @@ def __init__(self, binarySharedObject, diskPath, functionPointers, globalVariabl
 class BinarySharedObject:
     """Models a shared object library (.so) loadable on linux systems."""
 
-    def __init__(self, binaryForm, functionTypes, globalVariableDefinitions):
+    def __init__(self, binaryForm, functionTypes, globalVariableDefinitions, usedExternalFunctions):
         """
         Args:
             binaryForm - a bytes object containing the actual compiled code for the module
@@ -47,21 +47,27 @@ def __init__(self, binaryForm, functionTypes, globalVariableDefinitions):
         self.binaryForm = binaryForm
         self.functionTypes = functionTypes
         self.globalVariableDefinitions = globalVariableDefinitions
+        self.usedExternalFunctions = usedExternalFunctions
         self.hash = sha_hash(binaryForm)
 
     @property
     def definedSymbols(self):
         return self.functionTypes.keys()
 
     @staticmethod
-    def fromDisk(path, globalVariableDefinitions, functionNameToType):
+    def fromDisk(path, globalVariableDefinitions, functionNameToType, usedExternalFunctions):
         with open(path, "rb") as f:
             binaryForm = f.read()
 
-        return BinarySharedObject(binaryForm, functionNameToType, globalVariableDefinitions)
+        return BinarySharedObject(
+            binaryForm,
+            functionNameToType,
+            globalVariableDefinitions,
+            usedExternalFunctions
+        )
 
     @staticmethod
-    def fromModule(module, globalVariableDefinitions, functionNameToType):
+    def fromModule(module, globalVariableDefinitions, functionNameToType, usedExternalFunctions):
         target_triple = llvm.get_process_triple()
         target = llvm.Target.from_triple(target_triple)
         target_machine_shared_object = target.create_target_machine(reloc='pic', codemodel='default')
@@ -82,7 +88,12 @@ def fromModule(module, globalVariableDefinitions, functionNameToType):
             )
 
             with open(os.path.join(tf, "module.so"), "rb") as so_file:
-                return BinarySharedObject(so_file.read(), functionNameToType, globalVariableDefinitions)
+                return BinarySharedObject(
+                    so_file.read(),
+                    functionNameToType,
+                    globalVariableDefinitions,
+                    usedExternalFunctions
+                )
 
     def load(self, storageDir):
         """Instantiate this .so in temporary storage and return a dict from symbol -> integer function pointer"""

typed_python/compiler/native_compiler/compiler_cache.py

@@ -137,6 +137,10 @@ def loadModuleByHash(self, moduleHash, nameToTypedCallTarget, nameToNativeFuncti
 
             with open(os.path.join(targetDir, "submodules.dat"), "rb") as f:
                 submodules = SerializationContext().deserialize(f.read(), ListOf(str))
+
+            with open(os.path.join(targetDir, "linkDependencies.dat"), "rb") as f:
+                linkDependencies = SerializationContext().deserialize(f.read(), ListOf(str))
+
         except Exception:
             self.markModuleHashInvalid(moduleHash)
             return False
@@ -159,7 +163,8 @@ def loadModuleByHash(self, moduleHash, nameToTypedCallTarget, nameToNativeFuncti
         loaded = BinarySharedObject.fromDisk(
             modulePath,
             globalVarDefs,
-            functionNameToNativeType
+            functionNameToNativeType,
+            linkDependencies
         ).loadFromPath(modulePath)
 
         self.loadedModules[moduleHash] = loaded
@@ -173,7 +178,7 @@ def loadModuleByHash(self, moduleHash, nameToTypedCallTarget, nameToNativeFuncti
 
         return True
 
-    def addModule(self, binarySharedObject, nameToTypedCallTarget, linkDependencies):
+    def addModule(self, binarySharedObject, nameToTypedCallTarget):
         """Add new code to the compiler cache.
 
         Args:
@@ -185,10 +190,14 @@ def addModule(self, binarySharedObject, nameToTypedCallTarget, linkDependencies)
         """
         dependentHashes = set()
 
-        for name in linkDependencies:
+        for name in binarySharedObject.usedExternalFunctions:
             dependentHashes.add(self.symbolToLoadedModuleHash[name])
 
-        path, hashToUse = self.writeModuleToDisk(binarySharedObject, nameToTypedCallTarget, dependentHashes)
+        path, hashToUse = self.writeModuleToDisk(
+            binarySharedObject,
+            nameToTypedCallTarget,
+            dependentHashes
+        )
 
         self.loadedModules[hashToUse] = (
             binarySharedObject.loadFromPath(os.path.join(path, "module.so"))
@@ -291,6 +300,14 @@ def writeModuleToDisk(self, binarySharedObject, nameToTypedCallTarget, submodule
         with open(os.path.join(tempTargetDir, "submodules.dat"), "wb") as f:
             f.write(SerializationContext().serialize(ListOf(str)(submodules), ListOf(str)))
 
+        with open(os.path.join(tempTargetDir, "linkDependencies.dat"), "wb") as f:
+            f.write(
+                SerializationContext().serialize(
+                    ListOf(str)(binarySharedObject.usedExternalFunctions),
+                    ListOf(str)
+                )
+            )
+
         try:
             os.rename(tempTargetDir, targetDir)
         except IOError:

typed_python/compiler/native_compiler/module_definition.py

@@ -22,11 +22,20 @@ class ModuleDefinition:
         moduleText - a string containing the llvm IR for the module
         functionList - a list of the names of exported functions
         globalDefinitions - a dict from name to a GlobalDefinition
+        usedExternalFunctions - a set of symbols of functions that were compiled
+            in different modules that we depend on. These will all be in some module.
     """
     GET_GLOBAL_VARIABLES_NAME = ".get_global_variables"
 
-    def __init__(self, moduleText, functionNameToType, globalVariableDefinitions):
+    def __init__(
+        self,
+        moduleText,
+        functionNameToType,
+        globalVariableDefinitions,
+        usedExternalFunctions
+    ):
         self.moduleText = moduleText
         self.functionNameToType = functionNameToType
         self.globalVariableDefinitions = globalVariableDefinitions
         self.hash = sha_hash(moduleText)
+        self.usedExternalFunctions = usedExternalFunctions

typed_python/compiler/native_compiler/native_ast_to_llvm.py

@@ -568,7 +568,8 @@ def __init__(self,
                  builder,
                  arg_assignments,
                  output_type,
-                 external_function_references
+                 external_function_references,
+                 usedExternalFunctions
                  ):
         self.function = function
 
@@ -582,6 +583,7 @@ def __init__(self,
         self.arg_assignments = arg_assignments
         self.output_type = output_type
         self.external_function_references = external_function_references
+        self.usedExternalFunctions = usedExternalFunctions
         self.tags_initialized = {}
         self.stack_slots = {}
 
@@ -712,7 +714,9 @@ def namedCallTargetToLLVM(self, target):
             func = self.external_function_references[target.name]
         elif target.name in self.converter._externallyDefinedFunctionTypes:
             # this function is defined in a shared object that we've loaded from a prior
-            # invocation
+            # invocation. We don't have the code available locally, so we can't inline it.
+            self.usedExternalFunctions.add(target.name)
+
             if target.name not in self.external_function_references:
                 func_type = llvmlite.ir.FunctionType(
                     type_to_llvm_type(target.output_type),
@@ -1683,6 +1687,7 @@ def add_functions(self, names_to_definitions):
 
         globalDefinitions = {}
         globalDefinitionsLlvmValues = {}
+        usedExternalFunctions = set()
 
         while names_to_definitions:
             for name in sorted(names_to_definitions):
@@ -1712,7 +1717,8 @@ def add_functions(self, names_to_definitions):
                         builder,
                         arg_assignments,
                         definition.output_type,
-                        external_function_references
+                        external_function_references,
+                        usedExternalFunctions
                     )
 
                     func_converter.setup()
@@ -1759,7 +1765,8 @@ def add_functions(self, names_to_definitions):
         return ModuleDefinition(
             str(module),
             functionTypes,
-            globalDefinitions
+            globalDefinitions,
+            usedExternalFunctions
         )
 
     def defineGlobalMetadataAccessor(self, module, globalDefinitions, globalDefinitionsLlvmValues):

typed_python/compiler/native_compiler/native_compiler.py

@@ -114,8 +114,7 @@ def addFunctions(
         # map from str to native_ast.Function
         functionDefinitions,
         # map from str to the TypedCallTarget for any function that's actually typed
-        typedCallTargets,
-        externallyUsed
+        typedCallTargets
     ):
         """Add a collection of functions to the compiler.
 
@@ -131,10 +130,10 @@ def addFunctions(
             loadedModule.linkGlobalVariables()
         else:
             binary = self._buildSharedObject(functionDefinitions)
+
             self.compilerCache.addModule(
                 binary,
                 typedCallTargets,
-                externallyUsed
             )
 
     def functionPointerByName(self, linkName) -> NativeFunctionPointer:
@@ -206,6 +205,7 @@ def _buildSharedObject(self, functions):
             mod,
             module.globalVariableDefinitions,
             module.functionNameToType,
+            module.usedExternalFunctions
         )
 
     def _buildModule(self, functions):

typed_python/compiler/python_to_native_converter.py

@@ -108,20 +108,9 @@ def buildAndLinkNewModule(self):
             return
 
         # get a set of function names that we depend on
-        externallyUsed = set()
-
-        for funcName in targets:
-            ident = self._identity_for_link_name.get(funcName)
-            if ident is not None:
-                for dep in self._dependencies.getNamesDependedOn(ident):
-                    depLN = self._link_name_for_identity.get(dep)
-                    if depLN not in targets:
-                        externallyUsed.add(depLN)
-
         self.nativeCompiler.addFunctions(
             targets,
-            {name: self._targets[name] for name in targets if name in self._targets},
-            externallyUsed
+            {name: self._targets[name] for name in targets if name in self._targets}
         )
 
     def extract_new_function_definitions(self):