|
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 |
|
|
