""" Implement utils for supporting retargeting of dispatchers. WARNING: Features defined in this file are experimental. The API may change without notice. """ import abc import weakref from numba.core import errors class RetargetCache: """Cache for retargeted dispatchers. The cache uses the original dispatcher as the key. """ container_type = weakref.WeakKeyDictionary def __init__(self): self._cache = self.container_type() self._stat_hit = 0 self._stat_miss = 0 def save_cache(self, orig_disp, new_disp): """Save a dispatcher associated with the given key. """ self._cache[orig_disp] = new_disp def load_cache(self, orig_disp): """Load a dispatcher associated with the given key. """ out = self._cache.get(orig_disp) if out is None: self._stat_miss += 1 else: self._stat_hit += 1 return out def items(self): """Returns the contents of the cache. """ return self._cache.items() def stats(self): """Returns stats regarding cache hit/miss. """ return {'hit': self._stat_hit, 'miss': self._stat_miss} class BaseRetarget(abc.ABC): """Abstract base class for retargeting logic. """ @abc.abstractmethod def check_compatible(self, orig_disp): """Check that the retarget is compatible. This method does not return anything meaningful (e.g. None) Incompatibility is signalled via raising an exception. """ pass @abc.abstractmethod def retarget(self, orig_disp): """Retargets the given dispatcher and returns a new dispatcher-like callable. Or, returns the original dispatcher if the the target_backend will not change. """ pass class BasicRetarget(BaseRetarget): """A basic retargeting implementation for a single output target. This class has two abstract methods/properties that subclasses must define. - `output_target` must return output target name. - `compile_retarget` must define the logic to retarget the given dispatcher. By default, this class uses `RetargetCache` as the internal cache. This can be modified by overriding the `.cache_type` class attribute. """ cache_type = RetargetCache def __init__(self): self.cache = self.cache_type() @abc.abstractproperty def output_target(self) -> str: """Returns the output target name. See numba/tests/test_retargeting.py for example usage. """ pass @abc.abstractmethod def compile_retarget(self, orig_disp): """Returns the retargeted dispatcher. See numba/tests/test_retargeting.py for example usage. """ pass def check_compatible(self, orig_disp): """ This implementation checks that `self.output_target == orig_disp._required_target_backend` """ required_target = orig_disp._required_target_backend output_target = self.output_target if required_target is not None: if output_target != required_target: m = ("The output target does match the required target: " f"{output_target} != {required_target}.") raise errors.CompilerError(m) def retarget(self, orig_disp): """Apply retargeting to orig_disp. The retargeted dispatchers are cached for future use. """ cache = self.cache opts = orig_disp.targetoptions # Skip if the original dispatcher is targeting the same output target if opts.get('target_backend') == self.output_target: return orig_disp cached = cache.load_cache(orig_disp) # No cache? if cached is None: out = self.compile_retarget(orig_disp) cache.save_cache(orig_disp, out) else: out = cached return out