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 c33483e5e..baf3b9dc3 100644 --- a/numba_cuda/numba/cuda/typing/templates.py +++ b/numba_cuda/numba/cuda/typing/templates.py @@ -789,6 +789,36 @@ def _get_jit_decorator(self): return jit + def _call_overload_func(self, args, kws): + """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 + 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. + + 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. + """ + # 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) + 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. @@ -826,7 +856,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