343 lines
12 KiB
Python
343 lines
12 KiB
Python
|
import ctypes
|
||
|
|
from ctypes import POINTER, c_bool, c_char_p, c_uint8, c_uint64, c_size_t
|
||
|
|
|
||
|
|
from llvmlite.binding import ffi, targets
|
||
|
|
|
||
|
|
|
||
|
|
class _LinkElement(ctypes.Structure):
|
||
|
|
_fields_ = [("element_kind", c_uint8),
|
||
|
|
("value", c_char_p),
|
||
|
|
("value_len", c_size_t)]
|
||
|
|
|
||
|
|
|
||
|
|
class _SymbolAddress(ctypes.Structure):
|
||
|
|
_fields_ = [("name", c_char_p), ("address", c_uint64)]
|
||
|
|
|
||
|
|
|
||
|
|
class JITLibraryBuilder:
|
||
|
|
"""
|
||
|
|
Create a library for linking by OrcJIT
|
||
|
|
|
||
|
|
OrcJIT operates like a linker: a number of compilation units and
|
||
|
|
dependencies are collected together and linked into a single dynamic library
|
||
|
|
that can export functions to other libraries or to be consumed directly as
|
||
|
|
entry points into JITted code. The native OrcJIT has a lot of memory
|
||
|
|
management complications so this API is designed to work well with Python's
|
||
|
|
garbage collection.
|
||
|
|
|
||
|
|
The creation of a new library is a bit like a linker command line where
|
||
|
|
compilation units, mostly as LLVM IR, and previously constructed libraries
|
||
|
|
are linked together, then loaded into memory, and the addresses of exported
|
||
|
|
symbols are extracted. Any static initializers are run and the exported
|
||
|
|
addresses and a resource tracker is produced. As long as the resource
|
||
|
|
tracker is referenced somewhere in Python, the exported addresses will be
|
||
|
|
valid. Once the resource tracker is garbage collected, the static
|
||
|
|
destructors will run and library will be unloaded from memory.
|
||
|
|
"""
|
||
|
|
def __init__(self):
|
||
|
|
self.__entries = []
|
||
|
|
self.__exports = set()
|
||
|
|
self.__imports = {}
|
||
|
|
|
||
|
|
def add_ir(self, llvmir):
|
||
|
|
"""
|
||
|
|
Adds a compilation unit to the library using LLVM IR as the input
|
||
|
|
format.
|
||
|
|
|
||
|
|
This takes a string or an object that can be converted to a string,
|
||
|
|
including IRBuilder, that contains LLVM IR.
|
||
|
|
"""
|
||
|
|
self.__entries.append((0, str(llvmir).encode('utf-8')))
|
||
|
|
return self
|
||
|
|
|
||
|
|
def add_native_assembly(self, asm):
|
||
|
|
"""
|
||
|
|
Adds a compilation unit to the library using native assembly as the
|
||
|
|
input format.
|
||
|
|
|
||
|
|
This takes a string or an object that can be converted to a string that
|
||
|
|
contains native assembly, which will be
|
||
|
|
parsed by LLVM.
|
||
|
|
"""
|
||
|
|
self.__entries.append((1, str(asm).encode('utf-8')))
|
||
|
|
return self
|
||
|
|
|
||
|
|
def add_object_img(self, data):
|
||
|
|
"""
|
||
|
|
Adds a compilation unit to the library using pre-compiled object code.
|
||
|
|
|
||
|
|
This takes the bytes of the contents of an object artifact which will be
|
||
|
|
loaded by LLVM.
|
||
|
|
"""
|
||
|
|
self.__entries.append((2, bytes(data)))
|
||
|
|
return self
|
||
|
|
|
||
|
|
def add_object_file(self, file_path):
|
||
|
|
"""
|
||
|
|
Adds a compilation unit to the library using pre-compiled object file.
|
||
|
|
|
||
|
|
This takes a string or path-like object that references an object file
|
||
|
|
which will be loaded by LLVM.
|
||
|
|
"""
|
||
|
|
with open(file_path, "rb") as f:
|
||
|
|
self.__entries.append((2, f.read()))
|
||
|
|
return self
|
||
|
|
|
||
|
|
def add_jit_library(self, name):
|
||
|
|
"""
|
||
|
|
Adds an existing JIT library as prerequisite.
|
||
|
|
|
||
|
|
The name of the library must match the one provided in a previous link
|
||
|
|
command.
|
||
|
|
"""
|
||
|
|
self.__entries.append((3, str(name).encode('utf-8')))
|
||
|
|
return self
|
||
|
|
|
||
|
|
def add_current_process(self):
|
||
|
|
"""
|
||
|
|
Allows the JITted library to access symbols in the current binary.
|
||
|
|
|
||
|
|
That is, it allows exporting the current binary's symbols, including
|
||
|
|
loaded libraries, as imports to the JITted
|
||
|
|
library.
|
||
|
|
"""
|
||
|
|
self.__entries.append((3, b''))
|
||
|
|
return self
|
||
|
|
|
||
|
|
def import_symbol(self, name, address):
|
||
|
|
"""
|
||
|
|
Register the *address* of global symbol *name*. This will make
|
||
|
|
it usable (e.g. callable) from LLVM-compiled functions.
|
||
|
|
"""
|
||
|
|
self.__imports[str(name)] = c_uint64(address)
|
||
|
|
return self
|
||
|
|
|
||
|
|
def export_symbol(self, name):
|
||
|
|
"""
|
||
|
|
During linking, extract the address of a symbol that was defined in one
|
||
|
|
of the compilation units.
|
||
|
|
|
||
|
|
This allows getting symbols, functions or global variables, out of the
|
||
|
|
JIT linked library. The addresses will be
|
||
|
|
available when the link method is called.
|
||
|
|
"""
|
||
|
|
self.__exports.add(str(name))
|
||
|
|
return self
|
||
|
|
|
||
|
|
def link(self, lljit, library_name):
|
||
|
|
"""
|
||
|
|
Link all the current compilation units into a JITted library and extract
|
||
|
|
the address of exported symbols.
|
||
|
|
|
||
|
|
An instance of the OrcJIT instance must be provided and this will be the
|
||
|
|
scope that is used to find other JITted libraries that are dependencies
|
||
|
|
and also be the place where this library will be defined.
|
||
|
|
|
||
|
|
After linking, the method will return a resource tracker that keeps the
|
||
|
|
library alive. This tracker also knows the addresses of any exported
|
||
|
|
symbols that were requested.
|
||
|
|
|
||
|
|
The addresses will be valid as long as the resource tracker is
|
||
|
|
referenced.
|
||
|
|
|
||
|
|
When the resource tracker is destroyed, the library will be cleaned up,
|
||
|
|
however, the name of the library cannot be reused.
|
||
|
|
"""
|
||
|
|
assert not lljit.closed, "Cannot add to closed JIT"
|
||
|
|
encoded_library_name = str(library_name).encode('utf-8')
|
||
|
|
assert len(encoded_library_name) > 0, "Library cannot be empty"
|
||
|
|
elements = (_LinkElement * len(self.__entries))()
|
||
|
|
for idx, (kind, value) in enumerate(self.__entries):
|
||
|
|
elements[idx].element_kind = c_uint8(kind)
|
||
|
|
elements[idx].value = c_char_p(value)
|
||
|
|
elements[idx].value_len = c_size_t(len(value))
|
||
|
|
exports = (_SymbolAddress * len(self.__exports))()
|
||
|
|
for idx, name in enumerate(self.__exports):
|
||
|
|
exports[idx].name = name.encode('utf-8')
|
||
|
|
|
||
|
|
imports = (_SymbolAddress * len(self.__imports))()
|
||
|
|
for idx, (name, addr) in enumerate(self.__imports.items()):
|
||
|
|
imports[idx].name = name.encode('utf-8')
|
||
|
|
imports[idx].address = addr
|
||
|
|
|
||
|
|
with ffi.OutputString() as outerr:
|
||
|
|
tracker = lljit._capi.LLVMPY_LLJIT_Link(
|
||
|
|
lljit._ptr,
|
||
|
|
encoded_library_name,
|
||
|
|
elements,
|
||
|
|
len(self.__entries),
|
||
|
|
imports,
|
||
|
|
len(self.__imports),
|
||
|
|
exports,
|
||
|
|
len(self.__exports),
|
||
|
|
outerr)
|
||
|
|
if not tracker:
|
||
|
|
raise RuntimeError(str(outerr))
|
||
|
|
return ResourceTracker(tracker,
|
||
|
|
library_name,
|
||
|
|
{name: exports[idx].address
|
||
|
|
for idx, name in enumerate(self.__exports)})
|
||
|
|
|
||
|
|
|
||
|
|
class ResourceTracker(ffi.ObjectRef):
|
||
|
|
"""
|
||
|
|
A resource tracker is created for each loaded JIT library and keeps the
|
||
|
|
module alive.
|
||
|
|
|
||
|
|
OrcJIT supports unloading libraries that are no longer used. This resource
|
||
|
|
tracker should be stored in any object that reference functions or constants
|
||
|
|
for a JITted library. When all references to the resource tracker are
|
||
|
|
dropped, this will trigger LLVM to unload the library and destroy any
|
||
|
|
functions.
|
||
|
|
|
||
|
|
Failure to keep resource trackers while calling a function or accessing a
|
||
|
|
symbol can result in crashes or memory corruption.
|
||
|
|
|
||
|
|
LLVM internally tracks references between different libraries, so only
|
||
|
|
"leaf" libraries need to be tracked.
|
||
|
|
"""
|
||
|
|
def __init__(self, ptr, name, addresses):
|
||
|
|
self.__addresses = addresses
|
||
|
|
self.__name = name
|
||
|
|
ffi.ObjectRef.__init__(self, ptr)
|
||
|
|
|
||
|
|
def __getitem__(self, item):
|
||
|
|
"""
|
||
|
|
Get the address of an exported symbol as an integer
|
||
|
|
"""
|
||
|
|
return self.__addresses[item]
|
||
|
|
|
||
|
|
@property
|
||
|
|
def name(self):
|
||
|
|
return self.__name
|
||
|
|
|
||
|
|
def _dispose(self):
|
||
|
|
with ffi.OutputString() as outerr:
|
||
|
|
if self._capi.LLVMPY_LLJIT_Dylib_Tracker_Dispose(self, outerr):
|
||
|
|
raise RuntimeError(str(outerr))
|
||
|
|
|
||
|
|
|
||
|
|
class LLJIT(ffi.ObjectRef):
|
||
|
|
"""
|
||
|
|
A OrcJIT-based LLVM JIT engine that can compile and run LLVM IR as a
|
||
|
|
collection of JITted dynamic libraries
|
||
|
|
|
||
|
|
The C++ OrcJIT API has a lot of memory ownership patterns that do not work
|
||
|
|
with Python. This API attempts to provide ones that are safe at the expense
|
||
|
|
of some features. Each LLJIT instance is a collection of JIT-compiled
|
||
|
|
libraries. In the C++ API, there is a "main" library; this API does not
|
||
|
|
provide access to the main library. Use the JITLibraryBuilder to create a
|
||
|
|
new named library instead.
|
||
|
|
"""
|
||
|
|
def __init__(self, ptr):
|
||
|
|
self._td = None
|
||
|
|
ffi.ObjectRef.__init__(self, ptr)
|
||
|
|
|
||
|
|
def lookup(self, dylib, fn):
|
||
|
|
"""
|
||
|
|
Find a function in this dynamic library and construct a new tracking
|
||
|
|
object for it
|
||
|
|
|
||
|
|
If the library or function do not exist, an exception will occur.
|
||
|
|
|
||
|
|
Parameters
|
||
|
|
----------
|
||
|
|
dylib : str or None
|
||
|
|
the name of the library containing the symbol
|
||
|
|
fn : str
|
||
|
|
the name of the function to get
|
||
|
|
"""
|
||
|
|
assert not self.closed, "Cannot lookup in closed JIT"
|
||
|
|
address = ctypes.c_uint64()
|
||
|
|
with ffi.OutputString() as outerr:
|
||
|
|
tracker = ffi.lib.LLVMPY_LLJITLookup(self,
|
||
|
|
dylib.encode("utf-8"),
|
||
|
|
fn.encode("utf-8"),
|
||
|
|
ctypes.byref(address),
|
||
|
|
outerr)
|
||
|
|
if not tracker:
|
||
|
|
raise RuntimeError(str(outerr))
|
||
|
|
|
||
|
|
return ResourceTracker(tracker, dylib, {fn: address.value})
|
||
|
|
|
||
|
|
@property
|
||
|
|
def target_data(self):
|
||
|
|
"""
|
||
|
|
The TargetData for this LLJIT instance.
|
||
|
|
"""
|
||
|
|
if self._td is not None:
|
||
|
|
return self._td
|
||
|
|
ptr = ffi.lib.LLVMPY_LLJITGetDataLayout(self)
|
||
|
|
self._td = targets.TargetData(ptr)
|
||
|
|
self._td._owned = True
|
||
|
|
return self._td
|
||
|
|
|
||
|
|
def _dispose(self):
|
||
|
|
if self._td is not None:
|
||
|
|
self._td.detach()
|
||
|
|
self._capi.LLVMPY_LLJITDispose(self)
|
||
|
|
|
||
|
|
|
||
|
|
def create_lljit_compiler(target_machine=None, *,
|
||
|
|
use_jit_link=False,
|
||
|
|
suppress_errors=False):
|
||
|
|
"""
|
||
|
|
Create an LLJIT instance
|
||
|
|
"""
|
||
|
|
with ffi.OutputString() as outerr:
|
||
|
|
lljit = ffi.lib.LLVMPY_CreateLLJITCompiler(target_machine,
|
||
|
|
suppress_errors,
|
||
|
|
use_jit_link,
|
||
|
|
outerr)
|
||
|
|
if not lljit:
|
||
|
|
raise RuntimeError(str(outerr))
|
||
|
|
|
||
|
|
return LLJIT(lljit)
|
||
|
|
|
||
|
|
|
||
|
|
ffi.lib.LLVMPY_LLJITLookup.argtypes = [
|
||
|
|
ffi.LLVMOrcLLJITRef,
|
||
|
|
c_char_p,
|
||
|
|
c_char_p,
|
||
|
|
POINTER(c_uint64),
|
||
|
|
POINTER(c_char_p),
|
||
|
|
]
|
||
|
|
ffi.lib.LLVMPY_LLJITLookup.restype = ffi.LLVMOrcDylibTrackerRef
|
||
|
|
|
||
|
|
ffi.lib.LLVMPY_LLJITGetDataLayout.argtypes = [
|
||
|
|
ffi.LLVMOrcLLJITRef,
|
||
|
|
]
|
||
|
|
ffi.lib.LLVMPY_LLJITGetDataLayout.restype = ffi.LLVMTargetDataRef
|
||
|
|
|
||
|
|
ffi.lib.LLVMPY_CreateLLJITCompiler.argtypes = [
|
||
|
|
ffi.LLVMTargetMachineRef,
|
||
|
|
c_bool,
|
||
|
|
c_bool,
|
||
|
|
POINTER(c_char_p),
|
||
|
|
]
|
||
|
|
ffi.lib.LLVMPY_CreateLLJITCompiler.restype = ffi.LLVMOrcLLJITRef
|
||
|
|
|
||
|
|
ffi.lib.LLVMPY_LLJITDispose.argtypes = [
|
||
|
|
ffi.LLVMOrcLLJITRef,
|
||
|
|
]
|
||
|
|
|
||
|
|
|
||
|
|
ffi.lib.LLVMPY_LLJIT_Link.argtypes = [
|
||
|
|
ffi.LLVMOrcLLJITRef,
|
||
|
|
c_char_p,
|
||
|
|
POINTER(_LinkElement),
|
||
|
|
c_size_t,
|
||
|
|
POINTER(_SymbolAddress),
|
||
|
|
c_size_t,
|
||
|
|
POINTER(_SymbolAddress),
|
||
|
|
c_size_t,
|
||
|
|
POINTER(c_char_p)
|
||
|
|
]
|
||
|
|
ffi.lib.LLVMPY_LLJIT_Link.restype = ffi.LLVMOrcDylibTrackerRef
|
||
|
|
|
||
|
|
ffi.lib.LLVMPY_LLJIT_Dylib_Tracker_Dispose.argtypes = [
|
||
|
|
ffi.LLVMOrcDylibTrackerRef,
|
||
|
|
POINTER(c_char_p)
|
||
|
|
]
|
||
|
|
ffi.lib.LLVMPY_LLJIT_Dylib_Tracker_Dispose.restype = c_bool
|