Spaces:
Running
Running
| import platform | |
| from ctypes import (POINTER, c_char_p, c_bool, c_void_p, | |
| c_int, c_uint64, c_size_t, CFUNCTYPE, string_at, cast, | |
| py_object, Structure) | |
| from llvmlite.binding import ffi, targets, object_file | |
| # Just check these weren't optimized out of the DLL. | |
| ffi.lib.LLVMPY_LinkInMCJIT | |
| def create_mcjit_compiler(module, target_machine, use_lmm=None): | |
| """ | |
| Create a MCJIT ExecutionEngine from the given *module* and | |
| *target_machine*. | |
| *lmm* controls whether the llvmlite memory manager is used. If not supplied, | |
| the default choice for the platform will be used (``True`` on 64-bit ARM | |
| systems, ``False`` otherwise). | |
| """ | |
| if use_lmm is None: | |
| use_lmm = platform.machine() in ('arm64', 'aarch64') | |
| with ffi.OutputString() as outerr: | |
| engine = ffi.lib.LLVMPY_CreateMCJITCompiler( | |
| module, target_machine, use_lmm, outerr) | |
| if not engine: | |
| raise RuntimeError(str(outerr)) | |
| target_machine._owned = True | |
| return ExecutionEngine(engine, module=module) | |
| def check_jit_execution(): | |
| """ | |
| Check the system allows execution of in-memory JITted functions. | |
| An exception is raised otherwise. | |
| """ | |
| errno = ffi.lib.LLVMPY_TryAllocateExecutableMemory() | |
| if errno != 0: | |
| raise OSError(errno, | |
| "cannot allocate executable memory. " | |
| "This may be due to security restrictions on your " | |
| "system, such as SELinux or similar mechanisms." | |
| ) | |
| class ExecutionEngine(ffi.ObjectRef): | |
| """An ExecutionEngine owns all Modules associated with it. | |
| Deleting the engine will remove all associated modules. | |
| It is an error to delete the associated modules. | |
| """ | |
| _object_cache = None | |
| def __init__(self, ptr, module): | |
| """ | |
| Module ownership is transferred to the EE | |
| """ | |
| self._modules = set([module]) | |
| self._td = None | |
| module._owned = True | |
| ffi.ObjectRef.__init__(self, ptr) | |
| def get_function_address(self, name): | |
| """ | |
| Return the address of the function named *name* as an integer. | |
| It's a fatal error in LLVM if the symbol of *name* doesn't exist. | |
| """ | |
| return ffi.lib.LLVMPY_GetFunctionAddress(self, name.encode("ascii")) | |
| def get_global_value_address(self, name): | |
| """ | |
| Return the address of the global value named *name* as an integer. | |
| It's a fatal error in LLVM if the symbol of *name* doesn't exist. | |
| """ | |
| return ffi.lib.LLVMPY_GetGlobalValueAddress(self, name.encode("ascii")) | |
| def add_global_mapping(self, gv, addr): | |
| ffi.lib.LLVMPY_AddGlobalMapping(self, gv, addr) | |
| def add_module(self, module): | |
| """ | |
| Ownership of module is transferred to the execution engine | |
| """ | |
| if module in self._modules: | |
| raise KeyError("module already added to this engine") | |
| ffi.lib.LLVMPY_AddModule(self, module) | |
| module._owned = True | |
| self._modules.add(module) | |
| def finalize_object(self): | |
| """ | |
| Make sure all modules owned by the execution engine are fully processed | |
| and "usable" for execution. | |
| """ | |
| ffi.lib.LLVMPY_FinalizeObject(self) | |
| def run_static_constructors(self): | |
| """ | |
| Run static constructors which initialize module-level static objects. | |
| """ | |
| ffi.lib.LLVMPY_RunStaticConstructors(self) | |
| def run_static_destructors(self): | |
| """ | |
| Run static destructors which perform module-level cleanup of static | |
| resources. | |
| """ | |
| ffi.lib.LLVMPY_RunStaticDestructors(self) | |
| def remove_module(self, module): | |
| """ | |
| Ownership of module is returned | |
| """ | |
| with ffi.OutputString() as outerr: | |
| if ffi.lib.LLVMPY_RemoveModule(self, module, outerr): | |
| raise RuntimeError(str(outerr)) | |
| self._modules.remove(module) | |
| module._owned = False | |
| def target_data(self): | |
| """ | |
| The TargetData for this execution engine. | |
| """ | |
| if self._td is not None: | |
| return self._td | |
| ptr = ffi.lib.LLVMPY_GetExecutionEngineTargetData(self) | |
| self._td = targets.TargetData(ptr) | |
| self._td._owned = True | |
| return self._td | |
| def enable_jit_events(self): | |
| """ | |
| Enable JIT events for profiling of generated code. | |
| Return value indicates whether connection to profiling tool | |
| was successful. | |
| """ | |
| ret = ffi.lib.LLVMPY_EnableJITEvents(self) | |
| return ret | |
| def _find_module_ptr(self, module_ptr): | |
| """ | |
| Find the ModuleRef corresponding to the given pointer. | |
| """ | |
| ptr = cast(module_ptr, c_void_p).value | |
| for module in self._modules: | |
| if cast(module._ptr, c_void_p).value == ptr: | |
| return module | |
| return None | |
| def add_object_file(self, obj_file): | |
| """ | |
| Add object file to the jit. object_file can be instance of | |
| :class:ObjectFile or a string representing file system path | |
| """ | |
| if isinstance(obj_file, str): | |
| obj_file = object_file.ObjectFileRef.from_path(obj_file) | |
| ffi.lib.LLVMPY_MCJITAddObjectFile(self, obj_file) | |
| def set_object_cache(self, notify_func=None, getbuffer_func=None): | |
| """ | |
| Set the object cache "notifyObjectCompiled" and "getBuffer" | |
| callbacks to the given Python functions. | |
| """ | |
| self._object_cache_notify = notify_func | |
| self._object_cache_getbuffer = getbuffer_func | |
| # Lifetime of the object cache is managed by us. | |
| self._object_cache = _ObjectCacheRef(self) | |
| # Note this doesn't keep a reference to self, to avoid reference | |
| # cycles. | |
| ffi.lib.LLVMPY_SetObjectCache(self, self._object_cache) | |
| def _raw_object_cache_notify(self, data): | |
| """ | |
| Low-level notify hook. | |
| """ | |
| if self._object_cache_notify is None: | |
| return | |
| module_ptr = data.contents.module_ptr | |
| buf_ptr = data.contents.buf_ptr | |
| buf_len = data.contents.buf_len | |
| buf = string_at(buf_ptr, buf_len) | |
| module = self._find_module_ptr(module_ptr) | |
| if module is None: | |
| # The LLVM EE should only give notifications for modules | |
| # known by us. | |
| raise RuntimeError("object compilation notification " | |
| "for unknown module %s" % (module_ptr,)) | |
| self._object_cache_notify(module, buf) | |
| def _raw_object_cache_getbuffer(self, data): | |
| """ | |
| Low-level getbuffer hook. | |
| """ | |
| if self._object_cache_getbuffer is None: | |
| return | |
| module_ptr = data.contents.module_ptr | |
| module = self._find_module_ptr(module_ptr) | |
| if module is None: | |
| # The LLVM EE should only give notifications for modules | |
| # known by us. | |
| raise RuntimeError("object compilation notification " | |
| "for unknown module %s" % (module_ptr,)) | |
| buf = self._object_cache_getbuffer(module) | |
| if buf is not None: | |
| # Create a copy, which will be freed by the caller | |
| data[0].buf_ptr = ffi.lib.LLVMPY_CreateByteString(buf, len(buf)) | |
| data[0].buf_len = len(buf) | |
| def _dispose(self): | |
| # The modules will be cleaned up by the EE | |
| for mod in self._modules: | |
| mod.detach() | |
| if self._td is not None: | |
| self._td.detach() | |
| self._modules.clear() | |
| self._object_cache = None | |
| self._capi.LLVMPY_DisposeExecutionEngine(self) | |
| class _ObjectCacheRef(ffi.ObjectRef): | |
| """ | |
| Internal: an ObjectCache instance for use within an ExecutionEngine. | |
| """ | |
| def __init__(self, obj): | |
| ptr = ffi.lib.LLVMPY_CreateObjectCache(_notify_c_hook, | |
| _getbuffer_c_hook, | |
| obj) | |
| ffi.ObjectRef.__init__(self, ptr) | |
| def _dispose(self): | |
| self._capi.LLVMPY_DisposeObjectCache(self) | |
| # ============================================================================ | |
| # FFI | |
| ffi.lib.LLVMPY_CreateMCJITCompiler.argtypes = [ | |
| ffi.LLVMModuleRef, | |
| ffi.LLVMTargetMachineRef, | |
| c_bool, | |
| POINTER(c_char_p), | |
| ] | |
| ffi.lib.LLVMPY_CreateMCJITCompiler.restype = ffi.LLVMExecutionEngineRef | |
| ffi.lib.LLVMPY_RemoveModule.argtypes = [ | |
| ffi.LLVMExecutionEngineRef, | |
| ffi.LLVMModuleRef, | |
| POINTER(c_char_p), | |
| ] | |
| ffi.lib.LLVMPY_RemoveModule.restype = c_bool | |
| ffi.lib.LLVMPY_AddModule.argtypes = [ | |
| ffi.LLVMExecutionEngineRef, | |
| ffi.LLVMModuleRef | |
| ] | |
| ffi.lib.LLVMPY_AddGlobalMapping.argtypes = [ffi.LLVMExecutionEngineRef, | |
| ffi.LLVMValueRef, | |
| c_void_p] | |
| ffi.lib.LLVMPY_FinalizeObject.argtypes = [ffi.LLVMExecutionEngineRef] | |
| ffi.lib.LLVMPY_GetExecutionEngineTargetData.argtypes = [ | |
| ffi.LLVMExecutionEngineRef | |
| ] | |
| ffi.lib.LLVMPY_GetExecutionEngineTargetData.restype = ffi.LLVMTargetDataRef | |
| ffi.lib.LLVMPY_TryAllocateExecutableMemory.argtypes = [] | |
| ffi.lib.LLVMPY_TryAllocateExecutableMemory.restype = c_int | |
| ffi.lib.LLVMPY_GetFunctionAddress.argtypes = [ | |
| ffi.LLVMExecutionEngineRef, | |
| c_char_p | |
| ] | |
| ffi.lib.LLVMPY_GetFunctionAddress.restype = c_uint64 | |
| ffi.lib.LLVMPY_GetGlobalValueAddress.argtypes = [ | |
| ffi.LLVMExecutionEngineRef, | |
| c_char_p | |
| ] | |
| ffi.lib.LLVMPY_GetGlobalValueAddress.restype = c_uint64 | |
| ffi.lib.LLVMPY_MCJITAddObjectFile.argtypes = [ | |
| ffi.LLVMExecutionEngineRef, | |
| ffi.LLVMObjectFileRef | |
| ] | |
| class _ObjectCacheData(Structure): | |
| _fields_ = [ | |
| ('module_ptr', ffi.LLVMModuleRef), | |
| ('buf_ptr', c_void_p), | |
| ('buf_len', c_size_t), | |
| ] | |
| _ObjectCacheNotifyFunc = CFUNCTYPE(None, py_object, | |
| POINTER(_ObjectCacheData)) | |
| _ObjectCacheGetBufferFunc = CFUNCTYPE(None, py_object, | |
| POINTER(_ObjectCacheData)) | |
| # XXX The ctypes function wrappers are created at the top-level, otherwise | |
| # there are issues when creating CFUNCTYPEs in child processes on CentOS 5 | |
| # 32 bits. | |
| _notify_c_hook = _ObjectCacheNotifyFunc( | |
| ExecutionEngine._raw_object_cache_notify) | |
| _getbuffer_c_hook = _ObjectCacheGetBufferFunc( | |
| ExecutionEngine._raw_object_cache_getbuffer) | |
| ffi.lib.LLVMPY_CreateObjectCache.argtypes = [_ObjectCacheNotifyFunc, | |
| _ObjectCacheGetBufferFunc, | |
| py_object] | |
| ffi.lib.LLVMPY_CreateObjectCache.restype = ffi.LLVMObjectCacheRef | |
| ffi.lib.LLVMPY_DisposeObjectCache.argtypes = [ffi.LLVMObjectCacheRef] | |
| ffi.lib.LLVMPY_SetObjectCache.argtypes = [ffi.LLVMExecutionEngineRef, | |
| ffi.LLVMObjectCacheRef] | |
| ffi.lib.LLVMPY_CreateByteString.restype = c_void_p | |
| ffi.lib.LLVMPY_CreateByteString.argtypes = [c_void_p, c_size_t] | |