Skip to content
Permalink
 
 
Cannot retrieve contributors at this time
"""
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