From f2a1418e58181a987365eacbd57ce671f8706a94 Mon Sep 17 00:00:00 2001 From: Yevhenii Havrylko Date: Thu, 9 Jul 2026 23:08:46 +0000 Subject: [PATCH 1/2] Fix duplicate overload type resolution --- numba_cuda/numba/cuda/typing/templates.py | 43 ++++++++++++++++++++++- 1 file changed, 42 insertions(+), 1 deletion(-) diff --git a/numba_cuda/numba/cuda/typing/templates.py b/numba_cuda/numba/cuda/typing/templates.py index c33483e5e..b9872f061 100644 --- a/numba_cuda/numba/cuda/typing/templates.py +++ b/numba_cuda/numba/cuda/typing/templates.py @@ -789,6 +789,46 @@ def _get_jit_decorator(self): return jit + def _call_overload_func(self, args, kws): + """Invoke the overload function, memoizing on argument types only. + + The overload function receives only the argument *types* -- never the + compiler ``Flags`` on the ConfigStack -- so its result (the + implementation ``pyfunc``, or a ``(signature, pyfunc)`` tuple) depends + solely on ``args``/``kws`` and is independent of flags such as + ``inline`` or ``lto``. ``_impl_cache`` however keys on the active + flags, so the same call resolved under two different flag contexts + (e.g. an LTO kernel's type-inference stage vs. a force-inlined + device-function compile) would otherwise re-execute the -- potentially + very expensive -- overload body once per context. Memoizing here keeps + the flag-sensitive ``_impl_cache`` intact (flags still decide which + compiled artifact is built and stored) while guaranteeing the overload + body runs at most once per (typing-context, argument-type) set. + """ + # Per-subclass cache: guard on ``__dict__`` so a subclass never reuses + # a parent template's cache (which would alias distinct overloads that + # happen to share argument types). + cls = type(self) + if "_overload_result_cache" not in cls.__dict__: + cls._overload_result_cache = {} + cache = cls._overload_result_cache + + try: + key = (self.context, tuple(args), tuple(sorted(kws.items()))) + except TypeError: + # Unhashable/unsortable argument types or kwargs: do not memoize. + return self._overload_func(*args, **kws) + try: + return cache[key] + except KeyError: + pass + except TypeError: + return self._overload_func(*args, **kws) + + result = self._overload_func(*args, **kws) + cache[key] = result + return result + def _build_impl(self, cache_key, args, kws): """Build and cache the implementation. @@ -826,7 +866,7 @@ def _build_impl(self, cache_key, args, kws): # problems raise TypingError(str(e)) from e else: - ovf_result = self._overload_func(*args, **kws) + ovf_result = self._call_overload_func(args, kws) if ovf_result is None: # No implementation => fail typing @@ -946,6 +986,7 @@ def make_overload_template( key=func, _overload_func=staticmethod(overload_func), _impl_cache={}, + _overload_result_cache={}, _compiled_overloads={}, _jit_options=jit_options, _strict=strict, From d0d6b96a886968e9ee2f107ae76f138cef5601a8 Mon Sep 17 00:00:00 2001 From: Yevhenii Havrylko Date: Thu, 9 Jul 2026 23:42:30 +0000 Subject: [PATCH 2/2] Simplify + add tests --- .../numba/cuda/tests/cudapy/test_overload.py | 139 ++++++++++++++++++ numba_cuda/numba/cuda/typing/templates.py | 45 +++--- 2 files changed, 156 insertions(+), 28 deletions(-) diff --git a/numba_cuda/numba/cuda/tests/cudapy/test_overload.py b/numba_cuda/numba/cuda/tests/cudapy/test_overload.py index 2bba8f4e7..62c6e9aa0 100644 --- a/numba_cuda/numba/cuda/tests/cudapy/test_overload.py +++ b/numba_cuda/numba/cuda/tests/cudapy/test_overload.py @@ -398,5 +398,144 @@ def kernel(a, b, out): self.assertEqual(tuple(out), (6, 2)) +@skip_on_cudasim("Overloading not supported in cudasim") +class TestOverloadFuncCaching(CUDATestCase): + """The overload body must execute at most once per argument-type set. + + Numba resolves the same overloaded call under several ConfigStack flag + contexts during a single kernel compilation (type inference vs. a + force-inlined / LTO device-function compile). ``_impl_cache`` keys on those + flags, so the -- potentially very expensive -- overload body would otherwise + run once per context. ``_OverloadFunctionTemplate`` memoizes the overload + result (which depends only on the argument types, never on the flags) to + collapse those to a single execution. These tests pin that down. + """ + + @staticmethod + def _make_template(overload_func, inline="never"): + from numba.cuda.typing.templates import make_overload_template + + def target(x): + pass + + return make_overload_template( + target, overload_func, jit_options={}, strict=True, inline=inline + ) + + def test_overload_body_runs_once(self): + calls = [] + + def ol(x): + calls.append(x) + + def impl(x): + pass + + return impl + + template = self._make_template(ol)(None) + argty = types.int32 + + r1 = template._call_overload_func((argty,), {}) + r2 = template._call_overload_func((argty,), {}) + + self.assertEqual(len(calls), 1) + self.assertIs(r1, r2) + + def test_distinct_arg_types_run_again(self): + calls = [] + + def ol(x): + calls.append(x) + + def impl(x): + pass + + return impl + + template = self._make_template(ol)(None) + + template._call_overload_func((types.int32,), {}) + template._call_overload_func((types.int64,), {}) + + self.assertEqual(len(calls), 2) + + def test_kwargs_participate_in_key(self): + calls = [] + + def ol(x, flag=False): + calls.append((x, flag)) + + def impl(x, flag=False): + pass + + return impl + + template = self._make_template(ol)(None) + + template._call_overload_func((types.int32,), {}) + template._call_overload_func((types.int32,), {}) + template._call_overload_func((types.int32,), {"flag": True}) + + self.assertEqual(len(calls), 2) + + def test_cache_is_per_template(self): + calls_a = [] + calls_b = [] + + def ol_a(x): + calls_a.append(x) + + def impl(x): + pass + + return impl + + def ol_b(x): + calls_b.append(x) + + def impl(x): + pass + + return impl + + template_a = self._make_template(ol_a)(None) + template_b = self._make_template(ol_b)(None) + argty = types.int32 + + ra1 = template_a._call_overload_func((argty,), {}) + ra2 = template_a._call_overload_func((argty,), {}) + rb1 = template_b._call_overload_func((argty,), {}) + + # Same argument type, but each template keeps its own cache: the two + # distinct overloads must not alias one another. + self.assertEqual(len(calls_a), 1) + self.assertEqual(len(calls_b), 1) + self.assertIs(ra1, ra2) + self.assertIsNot(ra1, rb1) + + def test_cache_lives_on_template_class(self): + # Template instances are transient -- Numba creates a fresh one per + # resolution -- so the cache must live on the template *class* for a + # second instance to reuse the first's result. + calls = [] + + def ol(x): + calls.append(x) + + def impl(x): + pass + + return impl + + template_cls = self._make_template(ol) + argty = types.int32 + + template_cls(None)._call_overload_func((argty,), {}) + template_cls(None)._call_overload_func((argty,), {}) + + self.assertEqual(len(calls), 1) + + if __name__ == "__main__": unittest.main() diff --git a/numba_cuda/numba/cuda/typing/templates.py b/numba_cuda/numba/cuda/typing/templates.py index b9872f061..baf3b9dc3 100644 --- a/numba_cuda/numba/cuda/typing/templates.py +++ b/numba_cuda/numba/cuda/typing/templates.py @@ -790,7 +790,7 @@ def _get_jit_decorator(self): return jit def _call_overload_func(self, args, kws): - """Invoke the overload function, memoizing on argument types only. + """Invoke the overload function, memoizing on the argument types only. The overload function receives only the argument *types* -- never the compiler ``Flags`` on the ConfigStack -- so its result (the @@ -800,34 +800,24 @@ def _call_overload_func(self, args, kws): flags, so the same call resolved under two different flag contexts (e.g. an LTO kernel's type-inference stage vs. a force-inlined device-function compile) would otherwise re-execute the -- potentially - very expensive -- overload body once per context. Memoizing here keeps - the flag-sensitive ``_impl_cache`` intact (flags still decide which - compiled artifact is built and stored) while guaranteeing the overload - body runs at most once per (typing-context, argument-type) set. + very expensive -- overload body once per context. + + Wrapping ``_overload_func`` in a per-subclass ``functools.lru_cache`` + collapses those to a single execution per argument-type set, while + leaving the flag-sensitive ``_impl_cache`` intact (flags still decide + which compiled artifact is built and stored). Argument types are + already required to be hashable by ``_impl_cache``, so no additional + constraint is introduced. """ - # Per-subclass cache: guard on ``__dict__`` so a subclass never reuses - # a parent template's cache (which would alias distinct overloads that - # happen to share argument types). + # Bind the cache lazily and per-subclass: guard on ``__dict__`` so a + # subclass never reuses a parent template's cached function (which would + # alias distinct overloads that happen to share argument types). cls = type(self) - if "_overload_result_cache" not in cls.__dict__: - cls._overload_result_cache = {} - cache = cls._overload_result_cache - - try: - key = (self.context, tuple(args), tuple(sorted(kws.items()))) - except TypeError: - # Unhashable/unsortable argument types or kwargs: do not memoize. - return self._overload_func(*args, **kws) - try: - return cache[key] - except KeyError: - pass - except TypeError: - return self._overload_func(*args, **kws) - - result = self._overload_func(*args, **kws) - cache[key] = result - return result + cached = cls.__dict__.get("_cached_overload_func") + if cached is None: + cached = functools.lru_cache(maxsize=None)(cls._overload_func) + cls._cached_overload_func = cached + return cached(*args, **kws) def _build_impl(self, cache_key, args, kws): """Build and cache the implementation. @@ -986,7 +976,6 @@ def make_overload_template( key=func, _overload_func=staticmethod(overload_func), _impl_cache={}, - _overload_result_cache={}, _compiled_overloads={}, _jit_options=jit_options, _strict=strict,