Skip to content

Commit 0193409

Browse files
committed
Update program cache linker backend decision
1 parent be557d0 commit 0193409

2 files changed

Lines changed: 29 additions & 18 deletions

File tree

‎cuda_core/cuda/core/utils/_program_cache/_keys.py‎

Lines changed: 23 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -173,12 +173,11 @@ def _nvrtc_version() -> tuple[int, int]:
173173
def _linker_backend_and_version(use_driver: bool) -> tuple[str, str]:
174174
"""Return ``(backend, version)`` for the linker used on PTX inputs.
175175
176-
``use_driver`` is the result of ``_decide_nvjitlink_or_driver()`` and
177-
must be passed in so a single ``make_program_cache_key`` call shares
178-
one probe across :meth:`_LinkerBackend.validate`,
179-
:meth:`option_fingerprint`, and :meth:`hash_version_probe` (otherwise
180-
a transient probe flap could write inconsistent fields into the same
181-
key).
176+
``use_driver`` is the cached result of the Linker's ``_choose_backend()``
177+
decision and must be passed in so a single ``make_program_cache_key`` call
178+
shares one probe across :meth:`_LinkerBackend.validate`,
179+
:meth:`option_fingerprint`, and :meth:`hash_version_probe` (otherwise a
180+
transient probe flap could write inconsistent fields into the same key).
182181
183182
Raises any underlying probe exception. ``make_program_cache_key`` catches
184183
and mixes the exception's class name into the digest, so the same probe
@@ -187,9 +186,9 @@ def _linker_backend_and_version(use_driver: bool) -> tuple[str, str]:
187186
working probe (``_probe_failed`` label vs. ``driver``/``nvrtc``/...).
188187
189188
nvJitLink version lookup goes through ``sys.modules`` first so we hit the
190-
same module ``_decide_nvjitlink_or_driver()`` already loaded. That keeps
191-
fingerprinting aligned with whichever ``cuda.bindings.nvjitlink`` import
192-
path the linker actually uses.
189+
same module the Linker probe already loaded. That keeps fingerprinting
190+
aligned with whichever ``cuda.bindings.nvjitlink`` import path the linker
191+
actually uses.
193192
"""
194193
import sys
195194

@@ -483,9 +482,21 @@ def _decide_driver(self) -> bool | None:
483482
"""
484483
if self._cached_decision is _DECISION_UNSET:
485484
try:
486-
from cuda.core._linker import _decide_nvjitlink_or_driver
487-
488-
self._cached_decision = _decide_nvjitlink_or_driver()
485+
from cuda.core import _linker
486+
487+
try:
488+
driver_major = _linker.driver_version()[0]
489+
except _linker.CUDAError:
490+
driver_major = None
491+
self._cached_decision = (
492+
_linker._choose_backend(
493+
driver_major,
494+
_linker._probe_nvjitlink(),
495+
inputs_have_ltoir=False,
496+
lto_requested=False,
497+
)
498+
== "driver"
499+
)
489500
except Exception as exc:
490501
self._cached_decision = None
491502
self._cached_decision_exc = exc

‎cuda_core/tests/test_program_cache.py‎

Lines changed: 6 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -406,7 +406,7 @@ def test_make_program_cache_key_ptx_driver_ignored_fields_collapse(field, a, b,
406406
compiles differing only in these flags produce identical ObjectCode."""
407407
from cuda.core import _linker
408408

409-
monkeypatch.setattr(_linker, "_decide_nvjitlink_or_driver", lambda: True) # driver
409+
monkeypatch.setattr(_linker, "_choose_backend", lambda *_args, **_kwargs: "driver")
410410
k_a = _make_key(code=".version 7.0", code_type="ptx", options=_opts(**{field: a}))
411411
k_b = _make_key(code=".version 7.0", code_type="ptx", options=_opts(**{field: b}))
412412
assert k_a == k_b
@@ -430,7 +430,7 @@ def test_make_program_cache_key_ptx_ptxas_options_canonicalized(a, b, monkeypatc
430430
as equivalent so equivalent compiles don't miss the cache."""
431431
from cuda.core import _linker
432432

433-
monkeypatch.setattr(_linker, "_decide_nvjitlink_or_driver", lambda: False) # nvJitLink
433+
monkeypatch.setattr(_linker, "_choose_backend", lambda *_args, **_kwargs: "nvjitlink")
434434
k_a = _make_key(code=".version 7.0", code_type="ptx", options=_opts(ptxas_options=a))
435435
k_b = _make_key(code=".version 7.0", code_type="ptx", options=_opts(ptxas_options=b))
436436
assert k_a == k_b
@@ -440,7 +440,7 @@ def test_make_program_cache_key_ptx_driver_ignored_fields_still_matter_under_nvj
440440
"""nvJitLink does honour those fields; they must still differentiate keys there."""
441441
from cuda.core import _linker
442442

443-
monkeypatch.setattr(_linker, "_decide_nvjitlink_or_driver", lambda: False) # nvJitLink
443+
monkeypatch.setattr(_linker, "_choose_backend", lambda *_args, **_kwargs: "nvjitlink")
444444
k_a = _make_key(code=".version 7.0", code_type="ptx", options=_opts(ftz=True))
445445
k_b = _make_key(code=".version 7.0", code_type="ptx", options=_opts(ftz=False))
446446
assert k_a != k_b
@@ -600,7 +600,7 @@ def test_make_program_cache_key_ptx_rejects_driver_linker_unsupported(option_kw,
600600
to exactly mirror the driver-linker's own gate."""
601601
from cuda.core import _linker
602602

603-
monkeypatch.setattr(_linker, "_decide_nvjitlink_or_driver", lambda: True) # driver
603+
monkeypatch.setattr(_linker, "_choose_backend", lambda *_args, **_kwargs: "driver")
604604
with pytest.raises(ValueError, match="driver linker"):
605605
_make_key(code=".version 7.0", code_type="ptx", options=_opts(**option_kw))
606606

@@ -610,7 +610,7 @@ def test_make_program_cache_key_ptx_accepts_driver_linker_unsupported_with_nvjit
610610
rejected at key time."""
611611
from cuda.core import _linker
612612

613-
monkeypatch.setattr(_linker, "_decide_nvjitlink_or_driver", lambda: False) # nvJitLink
613+
monkeypatch.setattr(_linker, "_choose_backend", lambda *_args, **_kwargs: "nvjitlink")
614614
# Should not raise.
615615
_make_key(code=".version 7.0", code_type="ptx", options=_opts(time=True))
616616

@@ -1104,7 +1104,7 @@ def test_make_program_cache_key_driver_probe_failure_taints_ptx_under_cuLink(mon
11041104
def _broken():
11051105
raise RuntimeError("driver probe failed")
11061106

1107-
monkeypatch.setattr(_linker, "_decide_nvjitlink_or_driver", lambda: True)
1107+
monkeypatch.setattr(_linker, "_choose_backend", lambda *_args, **_kwargs: "driver")
11081108
k_ok = _make_key(code=".ptx", code_type="ptx")
11091109
monkeypatch.setattr(_program_cache._keys, "_driver_version", _broken)
11101110
k_broken1 = _make_key(code=".ptx", code_type="ptx")

0 commit comments

Comments
 (0)