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