Spaces:
Running
Running
File size: 12,198 Bytes
3bb804c |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 220 221 222 223 224 225 226 227 228 229 230 231 232 233 234 235 236 237 238 239 240 241 242 243 244 245 246 247 248 249 250 251 252 253 254 255 256 257 258 259 260 261 262 263 264 265 266 267 268 269 270 271 272 273 274 275 276 277 278 279 280 281 282 283 284 285 286 287 288 289 290 291 292 293 294 295 296 297 298 299 300 301 302 303 304 305 306 307 308 309 310 311 312 313 314 315 316 317 318 319 320 321 322 323 324 325 326 327 328 329 330 331 332 333 334 335 336 337 338 339 340 341 342 343 |
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
|