NativeCompiler/compiler_cache pass function definitions to the llvm compiler to allow for cross-module-inlining.

Author: braxtonmckee

typed_python/compiler/native_compiler/binary_shared_object.py

@@ -36,38 +36,49 @@ 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, usedExternalFunctions):
+    def __init__(
+        self,
+        binaryForm,
+        functionTypes,
+        globalVariableDefinitions,
+        usedExternalFunctions,
+        functionDefinitions,
+    ):
         """
         Args:
             binaryForm - a bytes object containing the actual compiled code for the module
             functionTypes - a map from linkerName to native_ast.Type.Function containing the native
                 signatures of the functions
             globalVariableDefinitions - a map from name to GlobalVariableDefinition
+            usedExternalFunctions - a list of symbols defined in other modules that we need
+            functionDefinitions - dict from symbol to native_ast.Function object
         """
         self.binaryForm = binaryForm
         self.functionTypes = functionTypes
         self.globalVariableDefinitions = globalVariableDefinitions
         self.usedExternalFunctions = usedExternalFunctions
+        self.functionDefinitions = functionDefinitions
         self.hash = sha_hash(binaryForm)
 
     @property
     def definedSymbols(self):
         return self.functionTypes.keys()
 
     @staticmethod
-    def fromDisk(path, globalVariableDefinitions, functionNameToType, usedExternalFunctions):
+    def fromDisk(path, globalVariableDefinitions, functionNameToType, usedExternalFunctions, functionDefinitions):
         with open(path, "rb") as f:
             binaryForm = f.read()
 
         return BinarySharedObject(
             binaryForm,
             functionNameToType,
             globalVariableDefinitions,
-            usedExternalFunctions
+            usedExternalFunctions,
+            functionDefinitions
         )
 
     @staticmethod
-    def fromModule(module, globalVariableDefinitions, functionNameToType, usedExternalFunctions):
+    def fromModule(module, globalVariableDefinitions, functionNameToType, usedExternalFunctions, functionDefinitions):
         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')
@@ -92,7 +103,8 @@ def fromModule(module, globalVariableDefinitions, functionNameToType, usedExtern
                     so_file.read(),
                     functionNameToType,
                     globalVariableDefinitions,
-                    usedExternalFunctions
+                    usedExternalFunctions,
+                    functionDefinitions
                 )
 
     def load(self, storageDir):

typed_python/compiler/native_compiler/compiler_cache.py

@@ -15,6 +15,8 @@
 import os
 import uuid
 import shutil
+from typed_python.compiler.native_compiler.native_ast import Function
+from typed_python.compiler.native_compiler.native_ast_analysis import extractNamedCallTargets
 from typed_python.compiler.native_compiler.loaded_module import LoadedModule
 from typed_python.compiler.native_compiler.binary_shared_object import BinarySharedObject
 
@@ -47,8 +49,9 @@ class CompilerCache:
     by making it possible to determine if a given function is in the cache by organizing
     the manifests by, say, function name.
     """
-    def __init__(self, cacheDir):
+    def __init__(self, cacheDir, checkModuleValidity=True):
         self.cacheDir = cacheDir
+        self.checkModuleValidity = checkModuleValidity
 
         ensureDirExists(cacheDir)
 
@@ -102,9 +105,10 @@ def loadForSymbol(self, symbol):
 
             nameToTypedCallTarget = {}
             nameToNativeFunctionType = {}
+            nameToDefinition = {}
 
-            if self.loadModuleByHash(moduleHash, nameToTypedCallTarget, nameToNativeFunctionType):
-                return nameToTypedCallTarget, nameToNativeFunctionType
+            if self.loadModuleByHash(moduleHash, nameToTypedCallTarget, nameToNativeFunctionType, nameToDefinition):
+                return nameToTypedCallTarget, nameToNativeFunctionType, nameToDefinition
             else:
                 assert (
                     # either we can't load this symbol at all anymore
@@ -113,7 +117,13 @@ def loadForSymbol(self, symbol):
                     or moduleHash not in self.symbolToModuleHashes[symbol]
                 )
 
-    def loadModuleByHash(self, moduleHash, nameToTypedCallTarget, nameToNativeFunctionType):
+    def loadModuleByHash(
+        self,
+        moduleHash,
+        nameToTypedCallTarget,
+        nameToNativeFunctionType,
+        nameToDefinition
+    ):
         """Load a module by name.
 
         As we load, place all the newly imported typed call targets into
@@ -141,6 +151,11 @@ def loadModuleByHash(self, moduleHash, nameToTypedCallTarget, nameToNativeFuncti
             with open(os.path.join(targetDir, "linkDependencies.dat"), "rb") as f:
                 linkDependencies = SerializationContext().deserialize(f.read(), ListOf(str))
 
+            with open(os.path.join(targetDir, "functionDefinitions.dat"), "rb") as f:
+                functionDefinitions = SerializationContext().deserialize(
+                    f.read(), Dict(str, Function)
+                )
+
         except Exception:
             self.markModuleHashInvalid(moduleHash)
             return False
@@ -154,7 +169,8 @@ def loadModuleByHash(self, moduleHash, nameToTypedCallTarget, nameToNativeFuncti
             if not self.loadModuleByHash(
                 submodule,
                 nameToTypedCallTarget,
-                nameToNativeFunctionType
+                nameToNativeFunctionType,
+                nameToDefinition
             ):
                 return False
 
@@ -164,13 +180,15 @@ def loadModuleByHash(self, moduleHash, nameToTypedCallTarget, nameToNativeFuncti
             modulePath,
             globalVarDefs,
             functionNameToNativeType,
-            linkDependencies
+            linkDependencies,
+            functionDefinitions
         ).loadFromPath(modulePath)
 
         self.loadedModules[moduleHash] = loaded
 
         nameToTypedCallTarget.update(callTargets)
         nameToNativeFunctionType.update(functionNameToNativeType)
+        nameToDefinition.update(functionDefinitions)
 
         for symbol in functionNameToNativeType:
             if symbol not in self.symbolToLoadedModuleHash:
@@ -188,6 +206,37 @@ def addModule(self, binarySharedObject, nameToTypedCallTarget):
                 the formal python types for all the objects
             linkDependencies - a set of linknames we depend on directly.
         """
+        if self.checkModuleValidity:
+            externals = extractNamedCallTargets(
+                binarySharedObject.functionDefinitions
+            )
+
+            statedNames = set(binarySharedObject.usedExternalFunctions)
+
+            expectedNames = set(e.name for e in externals if not e.external) - set(
+                binarySharedObject.functionDefinitions
+            )
+
+            if statedNames != expectedNames:
+                if expectedNames - statedNames:
+                    raise Exception(
+                        "Invalid shared object - link dependencies don't match "
+                        + "stated shared object dependencies:\n\n"
+                        + "".join(
+                            ['    ' + x + "\n" for x in sorted(expectedNames - statedNames)]
+                        )
+                        + "\nwere referenced but not claimed in the manifest."
+                    )
+                else:
+                    raise Exception(
+                        "Invalid shared object - link dependencies don't match "
+                        + "stated shared object dependencies:\n\n"
+                        + "".join(
+                            ['    ' + x + "\n" for x in sorted(statedNames - expectedNames)]
+                        )
+                        + "\nwere claimed in the manifest but don't seem to be referenced"
+                    )
+
         dependentHashes = set()
 
         for name in binarySharedObject.usedExternalFunctions:
@@ -308,6 +357,14 @@ def writeModuleToDisk(self, binarySharedObject, nameToTypedCallTarget, submodule
                 )
             )
 
+        with open(os.path.join(tempTargetDir, "functionDefinitions.dat"), "wb") as f:
+            f.write(
+                SerializationContext().serialize(
+                    Dict(str, Function)(binarySharedObject.functionDefinitions),
+                    Dict(str, Function)
+                )
+            )
+
         try:
             os.rename(tempTargetDir, targetDir)
         except IOError:
@@ -324,6 +381,6 @@ def function_pointer_by_name(self, linkName):
             raise Exception("Can't find a module for " + linkName)
 
         if moduleHash not in self.loadedModules:
-            self.loadForSymbol(linkName)
+            raise Exception("You need to call 'loadForSymbol' on this linkName first")
 
         return self.loadedModules[moduleHash].functionPointers[linkName]

typed_python/compiler/native_compiler/compiler_cache_test.py

@@ -50,6 +50,57 @@ def test_compiler_cache_can_handle_conflicting_versions_of_the_same_code():
         assert len(os.listdir(compilerCacheDir)) == 2
 
 
+@pytest.mark.skipif('sys.platform=="darwin"')
+def test_compiler_cache_inlining():
+    fListDef = "\n".join([
+        "from typed_python import Entrypoint",
+        "@Entrypoint",
+        "def fList(aList):",
+        "    res = 0",
+        "    for x in aList:",
+        "        res += f(x)",
+        "    return res",
+    ])
+    xmodule = "\n".join([
+        "def f(x):",
+        "    return x",
+        fListDef,
+    ])
+    ymodule = "\n".join([
+        "from x import f",
+        fListDef,
+    ])
+    testmodule = "\n".join([
+        "import time",
+        "from typed_python import Entrypoint, ListOf",
+        "from x import fList",
+        "@Entrypoint",
+        "def makeRange(N):",
+        "    res = ListOf(int)()",
+        "    for i in range(N):",
+        "        res.append(i)",
+        "    return res",
+        "N = 10000000",
+        "bigRange = makeRange(N)",
+        "fList(bigRange)",
+        "t0 = time.time()",
+        "fList(bigRange)",
+        "duration = time.time() - t0",
+    ])
+
+    VERSION1 = {'x.py': xmodule, 'testmodule.py': testmodule}
+    VERSION2 = {'x.py': xmodule, 'y.py': ymodule, 'testmodule.py': testmodule.replace('x', 'y')}
+
+    with tempfile.TemporaryDirectory() as compilerCacheDir:
+        dur1 = evaluateExprInFreshProcess(VERSION1, 'testmodule.duration', compilerCacheDir)
+        dur2 = evaluateExprInFreshProcess(VERSION2, 'testmodule.duration', compilerCacheDir)
+
+        print(dur1)
+        print(dur2)
+
+        assert dur2 < 2 * dur1
+
+
 @pytest.mark.skipif('sys.platform=="darwin"')
 def test_compiler_cache_can_detect_invalidation_through_modules():
     xmodule = "\n".join([

typed_python/compiler/native_compiler/module_definition.py

@@ -24,6 +24,7 @@ class ModuleDefinition:
         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.
+        functionDefinitions - a dict from name to Function
     """
     GET_GLOBAL_VARIABLES_NAME = ".get_global_variables"
 
@@ -32,10 +33,12 @@ def __init__(
         moduleText,
         functionNameToType,
         globalVariableDefinitions,
-        usedExternalFunctions
+        usedExternalFunctions,
+        functionDefinitions
     ):
         self.moduleText = moduleText
         self.functionNameToType = functionNameToType
         self.globalVariableDefinitions = globalVariableDefinitions
         self.hash = sha_hash(moduleText)
         self.usedExternalFunctions = usedExternalFunctions
+        self.functionDefinitions = functionDefinitions

typed_python/compiler/native_compiler/native_ast.py

@@ -569,7 +569,6 @@ def expr_could_throw(self):
     Binop={'op': BinaryOp, 'left': Expression, 'right': Expression},
     Unaryop={'op': UnaryOp, 'operand': Expression},
     Variable={'name': str},
-    Attribute={'left': Expression, 'attr': str},
     StructElementByIndex={'left': Expression, 'index': int},
     ElementPtr={'left': Expression, 'offsets': TupleOf(Expression)},
     Call={'target': CallTarget, 'args': TupleOf(Expression)},

typed_python/compiler/native_compiler/native_ast_analysis.py

@@ -0,0 +1,80 @@
+#   Copyright 2023 typed_python Authors
+#
+#   Licensed under the Apache License, Version 2.0 (the "License");
+#   you may not use this file except in compliance with the License.
+#   You may obtain a copy of the License at
+#
+#       http://www.apache.org/licenses/LICENSE-2.0
+#
+#   Unless required by applicable law or agreed to in writing, software
+#   distributed under the License is distributed on an "AS IS" BASIS,
+#   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+#   See the License for the specific language governing permissions and
+#   limitations under the License.
+
+
+from typed_python import TupleOf, ListOf, Tuple, Dict, ConstDict, NamedTuple
+
+from typed_python.compiler.native_compiler.native_ast import (
+    Constant,
+    Type,
+    UnaryOp,
+    BinaryOp,
+    NamedCallTarget,
+    Expression,
+    Teardown,
+    ExpressionIntermediate,
+    CallTarget,
+    FunctionBody,
+    Function,
+    GlobalVariableMetadata
+)
+
+
+def visitAstChildren(node, callback):
+    if not callback(node):
+        return
+
+    # don't look in these
+    if isinstance(node, (UnaryOp, BinaryOp, Constant, Type, NamedCallTarget, GlobalVariableMetadata)):
+        return
+
+    if isinstance(node, (int, float, str, bytes, bool, type(None))):
+        return
+
+    if isinstance(node, Function):
+        visitAstChildren(node.args, callback)
+        visitAstChildren(node.body, callback)
+        visitAstChildren(node.output_type, callback)
+        return
+
+    if isinstance(node, (Expression, Teardown, ExpressionIntermediate, CallTarget, FunctionBody)):
+        for name in node.ElementType.ElementNames:
+            visitAstChildren(getattr(node, name), callback)
+        return
+
+    if isinstance(node, (Dict, ConstDict, dict)):
+        for k, v in node.items():
+            visitAstChildren(k, callback)
+            visitAstChildren(v, callback)
+        return
+
+    if isinstance(node, (TupleOf, ListOf, tuple, list, Tuple, NamedTuple)):
+        for child in node:
+            visitAstChildren(child, callback)
+        return
+
+    raise Exception(f"Unexpected AST node of type {type(node).__name__}")
+
+
+def extractNamedCallTargets(ast):
+    targets = set()
+
+    def check(node):
+        if isinstance(node, NamedCallTarget):
+            targets.add(node)
+        return True
+
+    visitAstChildren(ast, check)
+
+    return targets