diff --git a/numba_cuda/numba/cuda/core/config.py b/numba_cuda/numba/cuda/core/config.py index 16b440af0..d1672a471 100644 --- a/numba_cuda/numba/cuda/core/config.py +++ b/numba_cuda/numba/cuda/core/config.py @@ -424,6 +424,12 @@ def avx_default(): # Disable jit for debugging DISABLE_JIT = _readenv("NUMBA_DISABLE_JIT", int, 0) + # Disable caching of failed type inferences. + # Use this to isolate problems due to the fail cache. + DISABLE_TYPEINFER_FAIL_CACHE = _readenv( + "NUMBA_DISABLE_TYPEINFER_FAIL_CACHE", int, 0 + ) + # choose parallel backend to use THREADING_LAYER_PRIORITY = _readenv( "NUMBA_THREADING_LAYER_PRIORITY", diff --git a/numba_cuda/numba/cuda/typing/context.py b/numba_cuda/numba/cuda/typing/context.py index c9d4c6832..446fa051d 100644 --- a/numba_cuda/numba/cuda/typing/context.py +++ b/numba_cuda/numba/cuda/typing/context.py @@ -11,6 +11,7 @@ from importlib.util import find_spec from numba.cuda.core import errors +from numba.cuda.core import config from numba.cuda.typeconv import Conversion, rules from numba.cuda.typing.typeof import typeof, Purpose from numba.cuda.typing import templates @@ -50,6 +51,8 @@ class CallStack(Sequence): def __init__(self): self._stack = [] self._lock = threading.RLock() + # fail_cache only lasts for the current compilation session + self._fail_cache = {} def __getitem__(self, index): """ @@ -63,17 +66,34 @@ def __len__(self): @contextlib.contextmanager def register(self, target, typeinfer, func_id, args): - # guard compiling the same function with the same signature - if self.match(func_id.func, args): - msg = "compiler re-entrant to the same function signature" - raise errors.NumbaRuntimeError(msg) - self._lock.acquire() - self._stack.append(CallFrame(target, typeinfer, func_id, args)) - try: + with contextlib.ExitStack() as undo: + # guard compiling the same function with the same signature + if self.match(func_id.func, args): + msg = "compiler re-entrant to the same function signature" + raise errors.NumbaRuntimeError(msg) + + # Acquire lock + undo.enter_context(self._lock) + + # Clear fail_cache at the start and end of a compilation session + def clear_fail_cache(*exc): + if config.DISABLE_TYPEINFER_FAIL_CACHE: + return # bypass + # Clear cache if stack is empty + if not self._stack: + self._fail_cache.clear() + + clear_fail_cache() + undo.push(clear_fail_cache) + + # Setup callframe + self._stack.append(CallFrame(target, typeinfer, func_id, args)) + + @undo.push + def undo_stack(*exc): + self._stack.pop() + yield - finally: - self._stack.pop() - self._lock.release() def finditer(self, py_func): """ @@ -102,6 +122,64 @@ def match(self, py_func, args): if frame.args == args: return frame + def lookup_resolve_cache(self, func, args, kws): + """Lookup resolution cache for the given function type and argument + types. + """ + if not self._stack or config.DISABLE_TYPEINFER_FAIL_CACHE: + # if callstack is empty, bypass fail_cache + return _ResolveCache() + + def normalize_dict(obj): + if isinstance(obj, dict): + return tuple(sorted(obj.items())) + return obj # already hashable as-is + + def hashable(obj): + try: + hash(obj) + except TypeError: + return False + else: + return True + + key = func, args, normalize_dict(kws) + if not hashable(key): + return _ResolveCache() + return self._fail_cache.setdefault(key, _ResolveCache()) + + +class _ResolveCache: + """ + A cache for function resolution result. + Currently only remembers failed attempts. + """ + + def __init__(self): + self._status = "unmarked" + self._exc = None + + def mark_error(self, exc): + """Mark the function resolution as failed with an exception.""" + self._status = "error" + self._exc = exc + + def mark_failed(self): + """Mark the function resolution as failed.""" + self._status = "failed" + + def replay_failure(self): + """Replay the failure if it has been marked as failed or error.""" + if self._status == "error": + raise self._exc + else: + assert self._status == "failed" + return None + + def has_failed_previously(self): + """Return True if the function resolution has failed previously.""" + return self._status in {"failed", "error"} + class CallFrame: """ @@ -193,6 +271,10 @@ def resolve_function_type(self, func, args, kws): Resolve function type *func* for argument types *args* and *kws*. A signature is returned. """ + cache = self.callstack.lookup_resolve_cache(func, args, kws) + if cache.has_failed_previously(): + return cache.replay_failure() + # Prefer user definition first try: res = self._resolve_user_function_type(func, args, kws) @@ -212,8 +294,12 @@ def resolve_function_type(self, func, args, kws): # Re-raise last_exception if no function type has been found if res is None and last_exception is not None: + cache.mark_error(last_exception) raise last_exception + if res is None: + cache.mark_failed() + return res def _resolve_builtin_function_type(self, func, args, kws):