def __init__(self, pyfunc, sig, locals, options, pipeline_class=compiler.Compiler): args, return_type = sig if return_type is None: raise TypeError("C callback needs an explicit return type") self.__name__ = pyfunc.__name__ self.__qualname__ = getattr(pyfunc, '__qualname__', self.__name__) self.__wrapped__ = pyfunc self._pyfunc = pyfunc self._sig = signature(return_type, *args) self._compiler = _CFuncCompiler(pyfunc, self._targetdescr, options, locals, pipeline_class=pipeline_class) self._wrapper_name = None self._wrapper_address = None self._cache = NullCache() self._cache_hits = 0
def __init__(self, py_func, locals={}, targetoptions={}, impl_kind='direct', pipeline_class=compiler.Compiler): """ Parameters ---------- py_func: function object to be compiled locals: dict, optional Mapping of local variable names to Numba types. Used to override the types deduced by the type inference engine. targetoptions: dict, optional Target-specific config options. impl_kind: str Select the compiler mode for `@jit` and `@generated_jit` pipeline_class: type numba.compiler.CompilerBase The compiler pipeline type. """ self.typingctx = self.targetdescr.typing_context self.targetctx = self.targetdescr.target_context pysig = utils.pysignature(py_func) arg_count = len(pysig.parameters) can_fallback = not targetoptions.get('nopython', False) _DispatcherBase.__init__(self, arg_count, py_func, pysig, can_fallback, exact_match_required=False) functools.update_wrapper(self, py_func) self.targetoptions = targetoptions self.locals = locals self._cache = NullCache() compiler_class = self._impl_kinds[impl_kind] self._impl_kind = impl_kind self._compiler = compiler_class(py_func, self.targetdescr, targetoptions, locals, pipeline_class) self._cache_hits = collections.Counter() self._cache_misses = collections.Counter() self._type = types.Dispatcher(self) self.typingctx.insert_global(self, self._type) # Remember target restriction self._required_target_backend = targetoptions.get('target_backend')
def enable_caching(self): self._cache = FunctionCache(self._pyfunc)
class CFunc(object): """ A compiled C callback, as created by the @cfunc decorator. """ _targetdescr = registry.cpu_target def __init__(self, pyfunc, sig, locals, options, pipeline_class=compiler.Compiler): args, return_type = sig if return_type is None: raise TypeError("C callback needs an explicit return type") self.__name__ = pyfunc.__name__ self.__qualname__ = getattr(pyfunc, '__qualname__', self.__name__) self.__wrapped__ = pyfunc self._pyfunc = pyfunc self._sig = signature(return_type, *args) self._compiler = _CFuncCompiler(pyfunc, self._targetdescr, options, locals, pipeline_class=pipeline_class) self._wrapper_name = None self._wrapper_address = None self._cache = NullCache() self._cache_hits = 0 def enable_caching(self): self._cache = FunctionCache(self._pyfunc) @global_compiler_lock def compile(self): # Try to load from cache cres = self._cache.load_overload(self._sig, self._targetdescr.target_context) if cres is None: cres = self._compile_uncached() self._cache.save_overload(self._sig, cres) else: self._cache_hits += 1 self._library = cres.library self._wrapper_name = cres.fndesc.llvm_cfunc_wrapper_name self._wrapper_address = self._library.get_pointer_to_function( self._wrapper_name) def _compile_uncached(self): sig = self._sig # Compile native function cres = self._compiler.compile(sig.args, sig.return_type) assert not cres.objectmode # disabled by compiler above fndesc = cres.fndesc # Compile C wrapper # Note we reuse the same library to allow inlining the Numba # function inside the wrapper. library = cres.library module = library.create_ir_module(fndesc.unique_name) context = cres.target_context ll_argtypes = [context.get_value_type(ty) for ty in sig.args] ll_return_type = context.get_value_type(sig.return_type) wrapty = ir.FunctionType(ll_return_type, ll_argtypes) wrapfn = module.add_function(wrapty, fndesc.llvm_cfunc_wrapper_name) builder = ir.IRBuilder(wrapfn.append_basic_block('entry')) self._build_c_wrapper(context, builder, cres, wrapfn.args) library.add_ir_module(module) library.finalize() return cres def _build_c_wrapper(self, context, builder, cres, c_args): sig = self._sig pyapi = context.get_python_api(builder) fnty = context.call_conv.get_function_type(sig.return_type, sig.args) fn = builder.module.add_function(fnty, cres.fndesc.llvm_func_name) # XXX no obvious way to freeze an environment status, out = context.call_conv.call_function( builder, fn, sig.return_type, sig.args, c_args) with builder.if_then(status.is_error, likely=False): # If (and only if) an error occurred, acquire the GIL # and use the interpreter to write out the exception. gil_state = pyapi.gil_ensure() context.call_conv.raise_error(builder, pyapi, status) cstr = context.insert_const_string(builder.module, repr(self)) strobj = pyapi.string_from_string(cstr) pyapi.err_write_unraisable(strobj) pyapi.decref(strobj) pyapi.gil_release(gil_state) builder.ret(out) @property def native_name(self): """ The process-wide symbol the C callback is exposed as. """ # Note from our point of view, the C callback is the wrapper around # the native function. return self._wrapper_name @property def address(self): """ The address of the C callback. """ return self._wrapper_address @utils.cached_property def cffi(self): """ A cffi function pointer representing the C callback. """ import cffi ffi = cffi.FFI() # cffi compares types by name, so using precise types would risk # spurious mismatches (such as "int32_t" vs. "int"). return ffi.cast("void *", self.address) @utils.cached_property def ctypes(self): """ A ctypes function object representing the C callback. """ ctypes_args = [to_ctypes(ty) for ty in self._sig.args] ctypes_restype = to_ctypes(self._sig.return_type) functype = ctypes.CFUNCTYPE(ctypes_restype, *ctypes_args) return functype(self.address) def inspect_llvm(self): """ Return the LLVM IR of the C callback definition. """ return self._library.get_llvm_str() @property def cache_hits(self): return self._cache_hits def __repr__(self): return "<Numba C callback %r>" % (self.__qualname__,)
class Dispatcher(_DispatcherBase): """ Implementation of user-facing dispatcher objects (i.e. created using the @jit decorator). This is an abstract base class. Subclasses should define the targetdescr class attribute. """ _fold_args = True _impl_kinds = { 'direct': _FunctionCompiler, 'generated': _GeneratedFunctionCompiler, } # A {uuid -> instance} mapping, for deserialization _memo = weakref.WeakValueDictionary() # hold refs to last N functions deserialized, retaining them in _memo # regardless of whether there is another reference _recent = collections.deque(maxlen=config.FUNCTION_CACHE_SIZE) __uuid = None __numba__ = 'py_func' def __init__(self, py_func, locals={}, targetoptions={}, impl_kind='direct', pipeline_class=compiler.Compiler): """ Parameters ---------- py_func: function object to be compiled locals: dict, optional Mapping of local variable names to Numba types. Used to override the types deduced by the type inference engine. targetoptions: dict, optional Target-specific config options. impl_kind: str Select the compiler mode for `@jit` and `@generated_jit` pipeline_class: type numba.compiler.CompilerBase The compiler pipeline type. """ self.typingctx = self.targetdescr.typing_context self.targetctx = self.targetdescr.target_context pysig = utils.pysignature(py_func) arg_count = len(pysig.parameters) can_fallback = not targetoptions.get('nopython', False) _DispatcherBase.__init__(self, arg_count, py_func, pysig, can_fallback, exact_match_required=False) functools.update_wrapper(self, py_func) self.targetoptions = targetoptions self.locals = locals self._cache = NullCache() compiler_class = self._impl_kinds[impl_kind] self._impl_kind = impl_kind self._compiler = compiler_class(py_func, self.targetdescr, targetoptions, locals, pipeline_class) self._cache_hits = collections.Counter() self._cache_misses = collections.Counter() self._type = types.Dispatcher(self) self.typingctx.insert_global(self, self._type) def dump(self, tab=''): print(f'{tab}DUMP {type(self).__name__}[{self.py_func.__name__}, type code={self._type._code}]') for cres in self.overloads.values(): cres.dump(tab = tab + ' ') print(f'{tab}END DUMP {type(self).__name__}[{self.py_func.__name__}]') @property def _numba_type_(self): return types.Dispatcher(self) def enable_caching(self): self._cache = FunctionCache(self.py_func) def __get__(self, obj, objtype=None): '''Allow a JIT function to be bound as a method to an object''' if obj is None: # Unbound method return self else: # Bound method return pytypes.MethodType(self, obj) def __reduce__(self): """ Reduce the instance for pickling. This will serialize the original function as well the compilation options and compiled signatures, but not the compiled code itself. """ if self._can_compile: sigs = [] else: sigs = [cr.signature for cr in self.overloads.values()] globs = self._compiler.get_globals_for_reduction() return (serialize._rebuild_reduction, (self.__class__, str(self._uuid), serialize._reduce_function(self.py_func, globs), self.locals, self.targetoptions, self._impl_kind, self._can_compile, sigs)) @classmethod def _rebuild(cls, uuid, func_reduced, locals, targetoptions, impl_kind, can_compile, sigs): """ Rebuild an Dispatcher instance after it was __reduce__'d. """ try: return cls._memo[uuid] except KeyError: pass py_func = serialize._rebuild_function(*func_reduced) self = cls(py_func, locals, targetoptions, impl_kind) # Make sure this deserialization will be merged with subsequent ones self._set_uuid(uuid) for sig in sigs: self.compile(sig) self._can_compile = can_compile return self @property def _uuid(self): """ An instance-specific UUID, to avoid multiple deserializations of a given instance. Note this is lazily-generated, for performance reasons. """ u = self.__uuid if u is None: u = str(uuid.uuid1()) self._set_uuid(u) return u def _set_uuid(self, u): assert self.__uuid is None self.__uuid = u self._memo[u] = self self._recent.append(self) @global_compiler_lock def compile(self, sig): if not self._can_compile: raise RuntimeError("compilation disabled") # Use counter to track recursion compilation depth with self._compiling_counter: args, return_type = sigutils.normalize_signature(sig) # Don't recompile if signature already exists existing = self.overloads.get(tuple(args)) if existing is not None: return existing.entry_point # Try to load from disk cache cres = self._cache.load_overload(sig, self.targetctx) if cres is not None: self._cache_hits[sig] += 1 # XXX fold this in add_overload()? (also see compiler.py) if not cres.objectmode and not cres.interpmode: self.targetctx.insert_user_function(cres.entry_point, cres.fndesc, [cres.library]) self.add_overload(cres) return cres.entry_point self._cache_misses[sig] += 1 try: cres = self._compiler.compile(args, return_type) except errors.ForceLiteralArg as e: def folded(args, kws): return self._compiler.fold_argument_types(args, kws)[1] raise e.bind_fold_arguments(folded) self.add_overload(cres) self._cache.save_overload(sig, cres) return cres.entry_point def get_compile_result(self, sig): """Compile (if needed) and return the compilation result with the given signature. """ atypes = tuple(sig.args) if atypes not in self.overloads: self.compile(atypes) return self.overloads[atypes] def recompile(self): """ Recompile all signatures afresh. """ sigs = list(self.overloads) old_can_compile = self._can_compile # Ensure the old overloads are disposed of, including compiled functions. self._make_finalizer()() self._reset_overloads() self._cache.flush() self._can_compile = True try: for sig in sigs: self.compile(sig) finally: self._can_compile = old_can_compile @property def stats(self): return _CompileStats( cache_path=self._cache.cache_path, cache_hits=self._cache_hits, cache_misses=self._cache_misses, ) def parallel_diagnostics(self, signature=None, level=1): """ Print parallel diagnostic information for the given signature. If no signature is present it is printed for all known signatures. level is used to adjust the verbosity, level=1 (default) is minimal verbosity, and 2, 3, and 4 provide increasing levels of verbosity. """ def dump(sig): ol = self.overloads[sig] pfdiag = ol.metadata.get('parfor_diagnostics', None) if pfdiag is None: msg = "No parfors diagnostic available, is 'parallel=True' set?" raise ValueError(msg) pfdiag.dump(level) if signature is not None: dump(signature) else: [dump(sig) for sig in self.signatures] def get_metadata(self, signature=None): """ Obtain the compilation metadata for a given signature. """ if signature is not None: return self.overloads[signature].metadata else: return dict((sig, self.overloads[sig].metadata) for sig in self.signatures) def get_function_type(self): """Return unique function type of dispatcher when possible, otherwise return None. A Dispatcher instance has unique function type when it contains exactly one compilation result and its compilation has been disabled (via its disable_compile method). """ if not self._can_compile and len(self.overloads) == 1: cres = tuple(self.overloads.values())[0] return types.FunctionType(cres.signature)
def enable_caching(self): self.cache = FunctionCache(self.py_func)
def __init__(self, py_func, locals={}, targetoptions={}): self.py_func = py_func self.overloads = utils.UniqueDict() self.targetoptions = targetoptions self.locals = locals self.cache = NullCache()
class UFuncDispatcher(serialize.ReduceMixin): """ An object handling compilation of various signatures for a ufunc. """ targetdescr = ufunc_target def __init__(self, py_func, locals={}, targetoptions={}): self.py_func = py_func self.overloads = utils.UniqueDict() self.targetoptions = targetoptions self.locals = locals self.cache = NullCache() def _reduce_states(self): """ NOTE: part of ReduceMixin protocol """ return dict( pyfunc=self.py_func, locals=self.locals, targetoptions=self.targetoptions, ) @classmethod def _rebuild(cls, pyfunc, locals, targetoptions): """ NOTE: part of ReduceMixin protocol """ return cls(py_func=pyfunc, locals=locals, targetoptions=targetoptions) def enable_caching(self): self.cache = FunctionCache(self.py_func) def compile(self, sig, locals={}, **targetoptions): locs = self.locals.copy() locs.update(locals) topt = self.targetoptions.copy() topt.update(targetoptions) flags = compiler.Flags() self.targetdescr.options.parse_as_flags(flags, topt) flags.no_cpython_wrapper = True flags.error_model = "numpy" # Disable loop lifting # The feature requires a real # python function flags.enable_looplift = False return self._compile_core(sig, flags, locals) def _compile_core(self, sig, flags, locals): """ Trigger the compiler on the core function or load a previously compiled version from the cache. Returns the CompileResult. """ typingctx = self.targetdescr.typing_context targetctx = self.targetdescr.target_context @contextmanager def store_overloads_on_success(): # use to ensure overloads are stored on success try: yield except Exception: raise else: exists = self.overloads.get(cres.signature) if exists is None: self.overloads[cres.signature] = cres # Use cache and compiler in a critical section with global_compiler_lock: with store_overloads_on_success(): # attempt look up of existing cres = self.cache.load_overload(sig, targetctx) if cres is not None: return cres # Compile args, return_type = sigutils.normalize_signature(sig) cres = compiler.compile_extra(typingctx, targetctx, self.py_func, args=args, return_type=return_type, flags=flags, locals=locals) # cache lookup failed before so safe to save self.cache.save_overload(sig, cres) return cres
class CFunc(object): """ A compiled C callback, as created by the @cfunc decorator. """ _targetdescr = registry.cpu_target def __init__(self, pyfunc, sig, locals, options, pipeline_class=compiler.Compiler): args, return_type = sig if return_type is None: raise TypeError("C callback needs an explicit return type") self.__name__ = pyfunc.__name__ self.__qualname__ = getattr(pyfunc, '__qualname__', self.__name__) self.__wrapped__ = pyfunc self._pyfunc = pyfunc self._sig = signature(return_type, *args) self._compiler = _CFuncCompiler(pyfunc, self._targetdescr, options, locals, pipeline_class=pipeline_class) self._wrapper_name = None self._wrapper_address = None self._cache = NullCache() self._cache_hits = 0 def enable_caching(self): self._cache = FunctionCache(self._pyfunc) @global_compiler_lock def compile(self): # Try to load from cache cres = self._cache.load_overload(self._sig, self._targetdescr.target_context) if cres is None: cres = self._compile_uncached() self._cache.save_overload(self._sig, cres) else: self._cache_hits += 1 self._library = cres.library self._wrapper_name = cres.fndesc.llvm_cfunc_wrapper_name self._wrapper_address = self._library.get_pointer_to_function( self._wrapper_name) def _compile_uncached(self): sig = self._sig # Compile native function as well as cfunc wrapper return self._compiler.compile(sig.args, sig.return_type) @property def native_name(self): """ The process-wide symbol the C callback is exposed as. """ # Note from our point of view, the C callback is the wrapper around # the native function. return self._wrapper_name @property def address(self): """ The address of the C callback. """ return self._wrapper_address @utils.cached_property def cffi(self): """ A cffi function pointer representing the C callback. """ import cffi ffi = cffi.FFI() # cffi compares types by name, so using precise types would risk # spurious mismatches (such as "int32_t" vs. "int"). return ffi.cast("void *", self.address) @utils.cached_property def ctypes(self): """ A ctypes function object representing the C callback. """ ctypes_args = [to_ctypes(ty) for ty in self._sig.args] ctypes_restype = to_ctypes(self._sig.return_type) functype = ctypes.CFUNCTYPE(ctypes_restype, *ctypes_args) return functype(self.address) def inspect_llvm(self): """ Return the LLVM IR of the C callback definition. """ return self._library.get_llvm_str() @property def cache_hits(self): return self._cache_hits def __repr__(self): return "<Numba C callback %r>" % (self.__qualname__, ) def __call__(self, *args, **kwargs): return self._pyfunc(*args, **kwargs)
class Dispatcher(serialize.ReduceMixin, _MemoMixin, _DispatcherBase): """ Implementation of user-facing dispatcher objects (i.e. created using the @jit decorator). This is an abstract base class. Subclasses should define the targetdescr class attribute. """ _fold_args = True _impl_kinds = { 'direct': _FunctionCompiler, 'generated': _GeneratedFunctionCompiler, } __numba__ = 'py_func' def __init__(self, py_func, locals={}, targetoptions={}, impl_kind='direct', pipeline_class=compiler.Compiler): """ Parameters ---------- py_func: function object to be compiled locals: dict, optional Mapping of local variable names to Numba types. Used to override the types deduced by the type inference engine. targetoptions: dict, optional Target-specific config options. impl_kind: str Select the compiler mode for `@jit` and `@generated_jit` pipeline_class: type numba.compiler.CompilerBase The compiler pipeline type. """ self.typingctx = self.targetdescr.typing_context self.targetctx = self.targetdescr.target_context pysig = utils.pysignature(py_func) arg_count = len(pysig.parameters) can_fallback = not targetoptions.get('nopython', False) _DispatcherBase.__init__(self, arg_count, py_func, pysig, can_fallback, exact_match_required=False) functools.update_wrapper(self, py_func) self.targetoptions = targetoptions self.locals = locals self._cache = NullCache() compiler_class = self._impl_kinds[impl_kind] self._impl_kind = impl_kind self._compiler = compiler_class(py_func, self.targetdescr, targetoptions, locals, pipeline_class) self._cache_hits = collections.Counter() self._cache_misses = collections.Counter() self._type = types.Dispatcher(self) self.typingctx.insert_global(self, self._type) # Remember target restriction self._required_target_backend = targetoptions.get('target_backend') def dump(self, tab=''): print(f'{tab}DUMP {type(self).__name__}[{self.py_func.__name__}' f', type code={self._type._code}]') for cres in self.overloads.values(): cres.dump(tab=tab + ' ') print(f'{tab}END DUMP {type(self).__name__}[{self.py_func.__name__}]') @property def _numba_type_(self): return types.Dispatcher(self) def enable_caching(self): self._cache = FunctionCache(self.py_func) def __get__(self, obj, objtype=None): '''Allow a JIT function to be bound as a method to an object''' if obj is None: # Unbound method return self else: # Bound method return pytypes.MethodType(self, obj) def _reduce_states(self): """ Reduce the instance for pickling. This will serialize the original function as well the compilation options and compiled signatures, but not the compiled code itself. NOTE: part of ReduceMixin protocol """ if self._can_compile: sigs = [] else: sigs = [cr.signature for cr in self.overloads.values()] return dict( uuid=str(self._uuid), py_func=self.py_func, locals=self.locals, targetoptions=self.targetoptions, impl_kind=self._impl_kind, can_compile=self._can_compile, sigs=sigs, ) @classmethod def _rebuild(cls, uuid, py_func, locals, targetoptions, impl_kind, can_compile, sigs): """ Rebuild an Dispatcher instance after it was __reduce__'d. NOTE: part of ReduceMixin protocol """ try: return cls._memo[uuid] except KeyError: pass self = cls(py_func, locals, targetoptions, impl_kind) # Make sure this deserialization will be merged with subsequent ones self._set_uuid(uuid) for sig in sigs: self.compile(sig) self._can_compile = can_compile return self def compile(self, sig): disp = self._get_dispatcher_for_current_target() if disp is not self: return disp.compile(sig) with ExitStack() as scope: cres = None def cb_compiler(dur): if cres is not None: self._callback_add_compiler_timer(dur, cres) def cb_llvm(dur): if cres is not None: self._callback_add_llvm_timer(dur, cres) scope.enter_context( ev.install_timer("numba:compiler_lock", cb_compiler)) scope.enter_context(ev.install_timer("numba:llvm_lock", cb_llvm)) scope.enter_context(global_compiler_lock) if not self._can_compile: raise RuntimeError("compilation disabled") # Use counter to track recursion compilation depth with self._compiling_counter: args, return_type = sigutils.normalize_signature(sig) # Don't recompile if signature already exists existing = self.overloads.get(tuple(args)) if existing is not None: return existing.entry_point # Try to load from disk cache cres = self._cache.load_overload(sig, self.targetctx) if cres is not None: self._cache_hits[sig] += 1 # XXX fold this in add_overload()? (also see compiler.py) if not cres.objectmode: self.targetctx.insert_user_function( cres.entry_point, cres.fndesc, [cres.library]) self.add_overload(cres) return cres.entry_point self._cache_misses[sig] += 1 ev_details = dict( dispatcher=self, args=args, return_type=return_type, ) with ev.trigger_event("numba:compile", data=ev_details): try: cres = self._compiler.compile(args, return_type) except errors.ForceLiteralArg as e: def folded(args, kws): return self._compiler.fold_argument_types( args, kws)[1] raise e.bind_fold_arguments(folded) self.add_overload(cres) self._cache.save_overload(sig, cres) return cres.entry_point def get_compile_result(self, sig): """Compile (if needed) and return the compilation result with the given signature. """ atypes = tuple(sig.args) if atypes not in self.overloads: self.compile(atypes) return self.overloads[atypes] def recompile(self): """ Recompile all signatures afresh. """ sigs = list(self.overloads) old_can_compile = self._can_compile # Ensure the old overloads are disposed of, # including compiled functions. self._make_finalizer()() self._reset_overloads() self._cache.flush() self._can_compile = True try: for sig in sigs: self.compile(sig) finally: self._can_compile = old_can_compile @property def stats(self): return _CompileStats( cache_path=self._cache.cache_path, cache_hits=self._cache_hits, cache_misses=self._cache_misses, ) def parallel_diagnostics(self, signature=None, level=1): """ Print parallel diagnostic information for the given signature. If no signature is present it is printed for all known signatures. level is used to adjust the verbosity, level=1 (default) is minimal verbosity, and 2, 3, and 4 provide increasing levels of verbosity. """ def dump(sig): ol = self.overloads[sig] pfdiag = ol.metadata.get('parfor_diagnostics', None) if pfdiag is None: msg = "No parfors diagnostic available, is 'parallel=True' set?" raise ValueError(msg) pfdiag.dump(level) if signature is not None: dump(signature) else: [dump(sig) for sig in self.signatures] def get_metadata(self, signature=None): """ Obtain the compilation metadata for a given signature. """ if signature is not None: return self.overloads[signature].metadata else: return dict( (sig, self.overloads[sig].metadata) for sig in self.signatures) def get_function_type(self): """Return unique function type of dispatcher when possible, otherwise return None. A Dispatcher instance has unique function type when it contains exactly one compilation result and its compilation has been disabled (via its disable_compile method). """ if not self._can_compile and len(self.overloads) == 1: cres = tuple(self.overloads.values())[0] return types.FunctionType(cres.signature) def _get_retarget_dispatcher(self): """Returns a dispatcher for the retarget request. """ # Check TLS target configuration tc = TargetConfigurationStack() retarget = tc.get() retarget.check_compatible(self) disp = retarget.retarget(self) return disp def _get_dispatcher_for_current_target(self): """Returns a dispatcher for the current target registered in `TargetConfigurationStack`. `self` is returned if no target is specified. """ tc = TargetConfigurationStack() if tc: return self._get_retarget_dispatcher() else: return self def _call_tls_target(self, *args, **kwargs): """This is called when the C dispatcher logic sees a retarget request. """ disp = self._get_retarget_dispatcher() # Call the new dispatcher return disp(*args, **kwargs)