from __future__ import print_function, division, absolute_import from llvmlite.llvmpy.core import Type, Builder, Constant import llvmlite.llvmpy.core as lc from numba import types, cgutils, config 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 = Type.function(pyobj, [pyobj, pyobj, pyobj]) wrapper = self.module.add_function(wrapty, name=wrapname) builder = Builder(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(lc.ICMP_EQ, parseok, Constant.null(parseok.type)) 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) status, retval = self.context.call_conv.call_function( builder, self.func, self.fndesc.restype, self.fndesc.argtypes, innerargs) # 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) 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))