227 lines
8.3 KiB
Python
227 lines
8.3 KiB
Python
from llvmlite.ir import Constant, IRBuilder
|
|
import llvmlite.ir
|
|
|
|
from numba.core import types, config, cgutils
|
|
|
|
|
|
class _ArgManager(object):
|
|
"""
|
|
A utility class to handle argument unboxing and cleanup
|
|
"""
|
|
def __init__(self, context, builder, api, env_manager, endblk, nargs):
|
|
self.context = context
|
|
self.builder = builder
|
|
self.api = api
|
|
self.env_manager = env_manager
|
|
self.arg_count = 0 # how many function arguments have been processed
|
|
self.cleanups = []
|
|
self.nextblk = endblk
|
|
|
|
def add_arg(self, obj, ty):
|
|
"""
|
|
Unbox argument and emit code that handles any error during unboxing.
|
|
Args are cleaned up in reverse order of the parameter list, and
|
|
cleanup begins as soon as unboxing of any argument fails. E.g. failure
|
|
on arg2 will result in control flow going through:
|
|
|
|
arg2.err -> arg1.err -> arg0.err -> arg.end (returns)
|
|
"""
|
|
# Unbox argument
|
|
native = self.api.to_native_value(ty, obj)
|
|
|
|
# If an error occurred, go to the cleanup block for
|
|
# the previous argument
|
|
with cgutils.if_unlikely(self.builder, native.is_error):
|
|
self.builder.branch(self.nextblk)
|
|
|
|
# Define the cleanup function for the argument
|
|
def cleanup_arg():
|
|
# Native value reflection
|
|
self.api.reflect_native_value(ty, native.value, self.env_manager)
|
|
|
|
# Native value cleanup
|
|
if native.cleanup is not None:
|
|
native.cleanup()
|
|
|
|
# NRT cleanup
|
|
# (happens after the native value cleanup as the latter
|
|
# may need the native value)
|
|
if self.context.enable_nrt:
|
|
self.context.nrt.decref(self.builder, ty, native.value)
|
|
|
|
self.cleanups.append(cleanup_arg)
|
|
|
|
# Write the on-error cleanup block for this argument
|
|
cleanupblk = self.builder.append_basic_block(
|
|
"arg%d.err" % self.arg_count)
|
|
with self.builder.goto_block(cleanupblk):
|
|
cleanup_arg()
|
|
# Go to next cleanup block
|
|
self.builder.branch(self.nextblk)
|
|
|
|
self.nextblk = cleanupblk
|
|
self.arg_count += 1
|
|
return native.value
|
|
|
|
def emit_cleanup(self):
|
|
"""
|
|
Emit the cleanup code after returning from the wrapped function.
|
|
"""
|
|
for dtor in self.cleanups:
|
|
dtor()
|
|
|
|
|
|
class _GilManager(object):
|
|
"""
|
|
A utility class to handle releasing the GIL and then re-acquiring it
|
|
again.
|
|
"""
|
|
|
|
def __init__(self, builder, api, argman):
|
|
self.builder = builder
|
|
self.api = api
|
|
self.argman = argman
|
|
self.thread_state = api.save_thread()
|
|
|
|
def emit_cleanup(self):
|
|
self.api.restore_thread(self.thread_state)
|
|
self.argman.emit_cleanup()
|
|
|
|
|
|
class PyCallWrapper(object):
|
|
def __init__(self, context, module, func, fndesc, env, call_helper,
|
|
release_gil):
|
|
self.context = context
|
|
self.module = module
|
|
self.func = func
|
|
self.fndesc = fndesc
|
|
self.env = env
|
|
self.release_gil = release_gil
|
|
|
|
def build(self):
|
|
wrapname = self.fndesc.llvm_cpython_wrapper_name
|
|
|
|
# This is the signature of PyCFunctionWithKeywords
|
|
# (see CPython's methodobject.h)
|
|
pyobj = self.context.get_argument_type(types.pyobject)
|
|
wrapty = llvmlite.ir.FunctionType(pyobj, [pyobj, pyobj, pyobj])
|
|
wrapper = llvmlite.ir.Function(self.module, wrapty, name=wrapname)
|
|
|
|
builder = IRBuilder(wrapper.append_basic_block('entry'))
|
|
|
|
# - `closure` will receive the `self` pointer stored in the
|
|
# PyCFunction object (see _dynfunc.c)
|
|
# - `args` and `kws` will receive the tuple and dict objects
|
|
# of positional and keyword arguments, respectively.
|
|
closure, args, kws = wrapper.args
|
|
closure.name = 'py_closure'
|
|
args.name = 'py_args'
|
|
kws.name = 'py_kws'
|
|
|
|
api = self.context.get_python_api(builder)
|
|
self.build_wrapper(api, builder, closure, args, kws)
|
|
|
|
return wrapper, api
|
|
|
|
def build_wrapper(self, api, builder, closure, args, kws):
|
|
nargs = len(self.fndesc.argtypes)
|
|
|
|
objs = [api.alloca_obj() for _ in range(nargs)]
|
|
parseok = api.unpack_tuple(args, self.fndesc.qualname,
|
|
nargs, nargs, *objs)
|
|
|
|
pred = builder.icmp_unsigned(
|
|
'==',
|
|
parseok,
|
|
Constant(parseok.type, None))
|
|
with cgutils.if_unlikely(builder, pred):
|
|
builder.ret(api.get_null_object())
|
|
|
|
# Block that returns after erroneous argument unboxing/cleanup
|
|
endblk = builder.append_basic_block("arg.end")
|
|
with builder.goto_block(endblk):
|
|
builder.ret(api.get_null_object())
|
|
|
|
# Get the Environment object
|
|
env_manager = self.get_env(api, builder)
|
|
|
|
cleanup_manager = _ArgManager(self.context, builder, api,
|
|
env_manager, endblk, nargs)
|
|
|
|
# Compute the arguments to the compiled Numba function.
|
|
innerargs = []
|
|
for obj, ty in zip(objs, self.fndesc.argtypes):
|
|
if isinstance(ty, types.Omitted):
|
|
# It's an omitted value => ignore dummy Python object
|
|
innerargs.append(None)
|
|
else:
|
|
val = cleanup_manager.add_arg(builder.load(obj), ty)
|
|
innerargs.append(val)
|
|
|
|
if self.release_gil:
|
|
cleanup_manager = _GilManager(builder, api, cleanup_manager)
|
|
|
|
# We elect to not inline the top level user function into the call
|
|
# wrapper, this incurs an overhead of a function call, however, it
|
|
# increases optimisation stability in that the optimised user function
|
|
# is what will actually be run and it is this function that all the
|
|
# inspection tools "see". Further, this makes optimisation "stable" in
|
|
# that calling the user function from e.g. C or from this wrapper will
|
|
# result in the same code executing, were inlining permitted this may
|
|
# not be the case as the inline could trigger additional optimisation
|
|
# as the function goes into the wrapper, this resulting in the executing
|
|
# instruction stream being different from that of the instruction stream
|
|
# present in the user function.
|
|
status, retval = self.context.call_conv.call_function(
|
|
builder, self.func, self.fndesc.restype, self.fndesc.argtypes,
|
|
innerargs, attrs=('noinline',))
|
|
# Do clean up
|
|
self.debug_print(builder, "# callwrapper: emit_cleanup")
|
|
cleanup_manager.emit_cleanup()
|
|
self.debug_print(builder, "# callwrapper: emit_cleanup end")
|
|
|
|
# Determine return status
|
|
with builder.if_then(status.is_ok, likely=True):
|
|
# Ok => return boxed Python value
|
|
with builder.if_then(status.is_none):
|
|
api.return_none()
|
|
|
|
retty = self._simplified_return_type()
|
|
obj = api.from_native_return(retty, retval, env_manager)
|
|
builder.ret(obj)
|
|
|
|
# Error out
|
|
self.context.call_conv.raise_error(builder, api, status)
|
|
builder.ret(api.get_null_object())
|
|
|
|
def get_env(self, api, builder):
|
|
"""Get the Environment object which is declared as a global
|
|
in the module of the wrapped function.
|
|
"""
|
|
envname = self.context.get_env_name(self.fndesc)
|
|
gvptr = self.context.declare_env_global(builder.module, envname)
|
|
envptr = builder.load(gvptr)
|
|
|
|
env_body = self.context.get_env_body(builder, envptr)
|
|
|
|
api.emit_environment_sentry(envptr, return_pyobject=True,
|
|
debug_msg=self.fndesc.env_name)
|
|
env_manager = api.get_env_manager(self.env, env_body, envptr)
|
|
return env_manager
|
|
|
|
def _simplified_return_type(self):
|
|
"""
|
|
The NPM callconv has already converted simplified optional types.
|
|
We can simply use the value type from it.
|
|
"""
|
|
restype = self.fndesc.restype
|
|
# Optional type
|
|
if isinstance(restype, types.Optional):
|
|
return restype.type
|
|
else:
|
|
return restype
|
|
|
|
def debug_print(self, builder, msg):
|
|
if config.DEBUG_JIT:
|
|
self.context.debug_print(builder, "DEBUGJIT: {0}".format(msg))
|