Encapsulate 'CompilerCache' entirely within NativeCompiler.

braxtonmckee · braxtonmckee · commit ddfa7387adfe · 2023-05-24T21:38:42.000Z
diff --git a/typed_python/compiler/native_compiler/llvm_execution_engine.py b/typed_python/compiler/native_compiler/llvm_execution_engine.py
@@ -0,0 +1,78 @@
+#   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.
+
+import llvmlite.binding as llvm
+import llvmlite.ir
+import typed_python.compiler.native_compiler.native_ast_to_llvm as native_ast_to_llvm
+
+import ctypes
+from typed_python import _types
+
+llvm.initialize()
+llvm.initialize_native_target()
+llvm.initialize_native_asmprinter()  # yes, even this one
+
+target_triple = llvm.get_process_triple()
+target = llvm.Target.from_triple(target_triple)
+target_machine = target.create_target_machine()
+target_machine_shared_object = target.create_target_machine(reloc='pic', codemodel='default')
+
+ctypes.CDLL(_types.__file__, mode=ctypes.RTLD_GLOBAL)
+
+
+pointer_size = (
+    llvmlite.ir.PointerType(llvmlite.ir.DoubleType())
+    .get_abi_size(target_machine.target_data)
+)
+
+assert pointer_size == native_ast_to_llvm.pointer_size
+
+
+def sizeof_native_type(native_type):
+    if native_type.matches.Void:
+        return 0
+
+    return (
+        native_ast_to_llvm.type_to_llvm_type(native_type)
+        .get_abi_size(target_machine.target_data)
+    )
+
+
+# there can be only one llvm engine alive at once.
+_engineCache = []
+
+
+def create_execution_engine(inlineThreshold):
+    if _engineCache:
+        return _engineCache[0]
+
+    pmb = llvm.create_pass_manager_builder()
+    pmb.opt_level = 3
+    pmb.size_level = 0
+    pmb.inlining_threshold = inlineThreshold
+    pmb.loop_vectorize = True
+    pmb.slp_vectorize = True
+
+    pass_manager = llvm.create_module_pass_manager()
+    pmb.populate(pass_manager)
+
+    target_machine.add_analysis_passes(pass_manager)
+
+    # And an execution engine with an empty backing module
+    backing_mod = llvm.parse_assembly("")
+    engine = llvm.create_mcjit_compiler(backing_mod, target_machine)
+
+    _engineCache.append((engine, pass_manager))
+
+    return engine, pass_manager
diff --git a/typed_python/compiler/native_compiler/native_compiler.py b/typed_python/compiler/native_compiler/native_compiler.py
@@ -1,4 +1,4 @@
-#   Copyright 2019 typed_python Authors
+#   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.
@@ -13,75 +13,14 @@
 #   limitations under the License.
 
 import llvmlite.binding as llvm
-import llvmlite.ir
 import typed_python.compiler.native_compiler.native_ast as native_ast
 import typed_python.compiler.native_compiler.native_ast_to_llvm as native_ast_to_llvm
-
+from typed_python.compiler.native_compiler.compiler_cache import CompilerCache
+from typed_python.compiler.native_compiler.llvm_execution_engine import create_execution_engine
 from typed_python.compiler.native_compiler.loaded_module import LoadedModule
 from typed_python.compiler.native_compiler.native_function_pointer import NativeFunctionPointer
 from typed_python.compiler.native_compiler.binary_shared_object import BinarySharedObject
 
-import ctypes
-from typed_python import _types
-
-llvm.initialize()
-llvm.initialize_native_target()
-llvm.initialize_native_asmprinter()  # yes, even this one
-
-target_triple = llvm.get_process_triple()
-target = llvm.Target.from_triple(target_triple)
-target_machine = target.create_target_machine()
-target_machine_shared_object = target.create_target_machine(reloc='pic', codemodel='default')
-
-ctypes.CDLL(_types.__file__, mode=ctypes.RTLD_GLOBAL)
-
-
-pointer_size = (
-    llvmlite.ir.PointerType(llvmlite.ir.DoubleType())
-    .get_abi_size(target_machine.target_data)
-)
-
-assert pointer_size == native_ast_to_llvm.pointer_size
-
-
-def sizeof_native_type(native_type):
-    if native_type.matches.Void:
-        return 0
-
-    return (
-        native_ast_to_llvm.type_to_llvm_type(native_type)
-        .get_abi_size(target_machine.target_data)
-    )
-
-
-# there can be only one llvm engine alive at once.
-_engineCache = []
-
-
-def create_execution_engine(inlineThreshold):
-    if _engineCache:
-        return _engineCache[0]
-
-    pmb = llvm.create_pass_manager_builder()
-    pmb.opt_level = 3
-    pmb.size_level = 0
-    pmb.inlining_threshold = inlineThreshold
-    pmb.loop_vectorize = True
-    pmb.slp_vectorize = True
-
-    pass_manager = llvm.create_module_pass_manager()
-    pmb.populate(pass_manager)
-
-    target_machine.add_analysis_passes(pass_manager)
-
-    # And an execution engine with an empty backing module
-    backing_mod = llvm.parse_assembly("")
-    engine = llvm.create_mcjit_compiler(backing_mod, target_machine)
-
-    _engineCache.append((engine, pass_manager))
-
-    return engine, pass_manager
-
 
 class NativeCompiler:
     """"Engine for compiling bundles of native_ast.Function objects into NativeFunctionPointers.
@@ -91,15 +30,23 @@ class NativeCompiler:
         * compiling functions into a runnable form using llvm
         * performing any runtime-based performance optimizations
         * maintaining the compiler cache
+
+    Note that this class is NOT threadsafe and clients are expected to serialize their
+    access through Runtime.
     """
     def __init__(self, inlineThreshold):
+        self.compilerCache = None
         self.engine, self.module_pass_manager = create_execution_engine(inlineThreshold)
         self.converter = native_ast_to_llvm.Converter()
         self.functions_by_name = {}
         self.inlineThreshold = inlineThreshold
         self.verbose = False
         self.optimize = True
 
+    def initializeCompilerCache(self, compilerCacheDir):
+        """Indicate that we should use a compiler cache from disk at 'compilerCacheDir'."""
+        self.compilerCache = CompilerCache(compilerCacheDir)
+
     def markExternal(self, functionNameToType):
         """Provide type signatures for a set of external functions."""
         self.converter.markExternal(functionNameToType)
@@ -110,7 +57,65 @@ def mark_converter_verbose(self):
     def mark_llvm_codegen_verbose(self):
         self.verbose = True
 
-    def buildSharedObject(self, functions):
+    def addFunctions(
+        self,
+        # map from str to native_ast.Function
+        functionDefinitions,
+        # map from str to the TypedCallTarget for any function that's actually typed
+        typedCallTargets,
+        externallyUsed
+    ):
+        """Add a collection of functions to the compiler.
+
+        Once a function has been added, we can request a NativeFunctionPointer for it.
+        """
+        if self.compilerCache is None:
+            loadedModule = self._buildModule(functionDefinitions)
+            loadedModule.linkGlobalVariables()
+        else:
+            binary = self._buildSharedObject(functionDefinitions)
+
+            self.compilerCache.addModule(
+                binary,
+                typedCallTargets,
+                externallyUsed
+            )
+
+    def functionPointerByName(self, linkerName) -> NativeFunctionPointer:
+        """Find a NativeFunctionPointer for a given link-time name.
+
+        Args:
+            linkerName (str) - the name of the compiled symbol we want.
+
+        Returns:
+            a NativeFunctionPointer or None if the function has never been defined.
+        """
+        if self.compilerCache is not None:
+            # the compiler cache has every shared object and can load them
+            return self.compilerCache.function_pointer_by_name(linkerName)
+
+        # the llvm compiler is just building shared objects, but the
+        # compiler cache has all the pointers.
+        return self.functions_by_name.get(linkerName)
+
+    def loadFromCache(self, linkName):
+        """Attempt to load a cached copy of 'linkName' and all reachable code.
+
+        If it isn't defined, or has already been defined, return None. If we're loading it
+        for the first time, return a pair
+
+            (typedCallTargets, nativeTypes)
+
+        where typedCallTargets is a map from linkName to TypedCallTarget, and nativeTypes is
+        a map from linkName to native_ast.Type.Function giving the native implementation type.
+
+        WARNING: this will return None if you already called 'functionPointerByName' on it
+        """
+        if self.compilerCache:
+            if self.compilerCache.hasSymbol(linkName):
+                return self.compilerCache.loadForSymbol(linkName)
+
+    def _buildSharedObject(self, functions):
         """Add native definitions and return a BinarySharedObject representing the compiled code."""
         module = self.converter.add_functions(functions)
 
@@ -135,10 +140,7 @@ def buildSharedObject(self, functions):
             module.functionNameToType,
         )
 
-    def function_pointer_by_name(self, name):
-        return self.functions_by_name.get(name)
-
-    def buildModule(self, functions):
+    def _buildModule(self, functions):
         """Compile a list of functions into a new module.
 
         Args:
diff --git a/typed_python/compiler/native_compiler/native_compiler_test.py b/typed_python/compiler/native_compiler/native_compiler_test.py
@@ -57,7 +57,7 @@ def test_global_variable_pointers():
     # is global.
     nativeCompiler = Runtime.singleton().native_compiler
 
-    loadedModule = nativeCompiler.buildModule(
+    loadedModule = nativeCompiler._buildModule(
         dict(
             _readGlobalVarFunc=readGlobalVarFunc,
             _returnGlobalVarPtrFunc=returnGlobalVarPtrFunc
@@ -111,7 +111,7 @@ def test_create_binary_shared_object():
 
     nativeCompiler = Runtime.singleton().native_compiler
 
-    bso = nativeCompiler.buildSharedObject(
+    bso = nativeCompiler._buildSharedObject(
         {'__test_f_2': f}
     )
 
diff --git a/typed_python/compiler/python_to_native_converter.py b/typed_python/compiler/python_to_native_converter.py
@@ -41,11 +41,8 @@
 
 
 class PythonToNativeConverter:
-    def __init__(self, llvmCompiler, compilerCache):
-        object.__init__(self)
-
-        self.llvmCompiler = llvmCompiler
-        self.compilerCache = compilerCache
+    def __init__(self, nativeCompiler):
+        self.nativeCompiler = nativeCompiler
 
         # if True, then insert additional code to check for undefined behavior.
         self.generateDebugChecks = False
@@ -116,11 +113,6 @@ def buildAndLinkNewModule(self):
         if not targets:
             return
 
-        if self.compilerCache is None:
-            loadedModule = self.llvmCompiler.buildModule(targets)
-            loadedModule.linkGlobalVariables()
-            return
-
         # get a set of function names that we depend on
         externallyUsed = set()
 
@@ -132,10 +124,8 @@ def buildAndLinkNewModule(self):
                     if depLN not in targets:
                         externallyUsed.add(depLN)
 
-        binary = self.llvmCompiler.buildSharedObject(targets)
-
-        self.compilerCache.addModule(
-            binary,
+        self.nativeCompiler.addFunctions(
+            targets,
             {name: self._targets[name] for name in targets if name in self._targets},
             externallyUsed
         )
@@ -209,18 +199,16 @@ def defineLinkName(self, identity, linkName):
         return True
 
     def _loadFromCompilerCache(self, linkName):
-        if self.compilerCache:
-            if self.compilerCache.hasSymbol(linkName):
-                callTargetsAndTypes = self.compilerCache.loadForSymbol(linkName)
+        callTargetsAndTypes = self.nativeCompiler.loadFromCache(linkName)
 
-                if callTargetsAndTypes is not None:
-                    newTypedCallTargets, newNativeFunctionTypes = callTargetsAndTypes
+        if callTargetsAndTypes is not None:
+            newTypedCallTargets, newNativeFunctionTypes = callTargetsAndTypes
 
-                    self._targets.update(newTypedCallTargets)
-                    self.llvmCompiler.markExternal(newNativeFunctionTypes)
+            self._targets.update(newTypedCallTargets)
+            self.nativeCompiler.markExternal(newNativeFunctionTypes)
 
-                    self._allDefinedNames.update(newNativeFunctionTypes)
-                    self._allCachedNames.update(newNativeFunctionTypes)
+            self._allDefinedNames.update(newNativeFunctionTypes)
+            self._allCachedNames.update(newNativeFunctionTypes)
 
     def defineNonPythonFunction(self, name, identityTuple, context):
         """Define a non-python generating function (if we haven't defined it before already)
@@ -604,13 +592,7 @@ def functionPointerByName(self, linkerName) -> NativeFunctionPointer:
         Returns:
             a NativeFunctionPointer or None
         """
-        if self.compilerCache is None:
-            # the llvm compiler holds it all
-            return self.llvmCompiler.function_pointer_by_name(linkerName)
-        else:
-            # the llvm compiler is just building shared objects, but the
-            # compiler cache has all the pointers.
-            return self.compilerCache.function_pointer_by_name(linkerName)
+        return self.nativeCompiler.functionPointerByName(linkerName)
 
     def convertTypedFunctionCall(self, functionType, overloadIx, inputWrappers, assertIsRoot=False):
         overload = functionType.overloads[overloadIx]
diff --git a/typed_python/compiler/runtime.py b/typed_python/compiler/runtime.py
@@ -21,7 +21,6 @@
 import typed_python
 from typed_python.compiler.runtime_lock import runtimeLock
 from typed_python.compiler.conversion_level import ConversionLevel
-from typed_python.compiler.native_compiler.compiler_cache import CompilerCache
 from typed_python.type_function import TypeFunction
 from typed_python.compiler.type_wrappers.typed_tuple_masquerading_as_tuple_wrapper import TypedTupleMasqueradingAsTuple
 from typed_python.compiler.type_wrappers.named_tuple_masquerading_as_dict_wrapper import NamedTupleMasqueradingAsDict
@@ -201,16 +200,16 @@ def singleton():
         return _singleton[0]
 
     def __init__(self):
+        self.native_compiler = native_compiler.NativeCompiler(
+            inlineThreshold=100
+        )
         if os.getenv("TP_COMPILER_CACHE"):
-            self.compilerCache = CompilerCache(
+            self.native_compiler.initializeCompilerCache(
                 os.path.abspath(os.getenv("TP_COMPILER_CACHE"))
             )
-        else:
-            self.compilerCache = None
-        self.native_compiler = native_compiler.NativeCompiler(inlineThreshold=100)
+
         self.converter = python_to_native_converter.PythonToNativeConverter(
-            self.native_compiler,
-            self.compilerCache
+            self.native_compiler
         )
         self.lock = runtimeLock
         self.timesCompiled = 0
