Check nvJitLink driver major compatibility (#2320)

* Check nvJitLink driver major compatibility

* Narrow nvJitLink driver compatibility fallback

* Fix cuLink state storage lifetimes

* Isolate nvJitLink-specific cache tests

* Use public linker backend query in tests

* Use direct backend skips in program cache tests

---------

Co-authored-by: Michael Wang <isVoid@users.noreply.github.com>

Author: isVoid

cuda_core/cuda/core/_linker.pxd

@@ -2,15 +2,23 @@
 #
 # SPDX-License-Identifier: Apache-2.0
 
+from libcpp.vector cimport vector
+
+from cuda.bindings cimport cydriver
+
 from ._resource_handles cimport NvJitLinkHandle, CuLinkHandle
 
 
 cdef class Linker:
     cdef:
         NvJitLinkHandle _nvjitlink_handle
         CuLinkHandle _culink_handle
+        # The driver retains these arrays until cuLinkDestroy. Declare them
+        # after the handle so their destructors run after cuLinkDestroy.
+        vector[cydriver.CUjit_option] _drv_jit_keys
+        vector[void*] _drv_jit_values
         bint _use_nvjitlink
-        object _drv_log_bufs  # formatted_options list (driver); None for nvjitlink; cleared in link()
+        object _drv_log_bufs  # formatted_options list (driver); None for nvjitlink
         str _info_log         # decoded log; None until link() or pre-link get_*_log()
         str _error_log        # decoded log; None until link() or pre-link get_*_log()
         object _options       # LinkerOptions

cuda_core/cuda/core/_linker.pyi

@@ -112,7 +112,8 @@ class LinkerOptions:
 
     Since the linker may choose to use nvJitLink or the driver APIs as the linking backend,
     not all options are applicable. When the system's installed nvJitLink is too old (<12.3),
-    or not installed, the driver APIs (cuLink) will be used instead.
+    not installed, or older than the CUDA driver major version, the driver APIs (cuLink)
+    will be used instead.
 
     Attributes
     ----------

cuda_core/cuda/core/_linker.pyx

@@ -39,6 +39,7 @@ from cuda.core._utils.cuda_utils import (
     driver,
     is_sequence,
 )
+from cuda.core._utils.version import driver_version
 from cuda.core.typing import CompilerBackendType, ObjectCodeFormatType
 
 if TYPE_CHECKING:
@@ -52,7 +53,6 @@ _keep_driver_in_stub: "cuda.bindings.driver.CUlinkState"
 _keep_nvjitlink_in_stub: "cuda.bindings.nvjitlink.nvJitLinkHandle"
 
 ctypedef const char* const_char_ptr
-ctypedef void* void_ptr
 
 __all__ = ["Linker", "LinkerOptions"]
 
@@ -155,10 +155,21 @@ cdef class Linker:
 
     def close(self) -> None:
         """Destroy this linker."""
+        cdef vector[cydriver.CUjit_option] empty_keys
+        cdef vector[void*] empty_values
         if self._use_nvjitlink:
             self._nvjitlink_handle.reset()
         else:
+            if self._drv_log_bufs is not None:
+                if self._info_log is None:
+                    self._info_log = self.get_info_log()
+                if self._error_log is None:
+                    self._error_log = self.get_error_log()
+            # Destroy the CUlinkState before releasing storage referenced by it.
             self._culink_handle.reset()
+            self._drv_jit_keys.swap(empty_keys)
+            self._drv_jit_values.swap(empty_values)
+            self._drv_log_bufs = None
 
     @property
     def handle(self) -> LinkerHandleT:
@@ -207,7 +218,8 @@ class LinkerOptions:
 
     Since the linker may choose to use nvJitLink or the driver APIs as the linking backend,
     not all options are applicable. When the system's installed nvJitLink is too old (<12.3),
-    or not installed, the driver APIs (cuLink) will be used instead.
+    not installed, or older than the CUDA driver major version, the driver APIs (cuLink)
+    will be used instead.
 
     Attributes
     ----------
@@ -473,8 +485,8 @@ cdef inline int Linker_init(Linker self, tuple object_codes, object options) exc
     cdef cydriver.CUlinkState c_raw_culink
     cdef Py_ssize_t c_num_opts, i
     cdef vector[const_char_ptr] c_str_opts
-    cdef vector[cydriver.CUjit_option] c_jit_keys
-    cdef vector[void_ptr] c_jit_values
+    cdef cydriver.CUjit_option* c_drv_jit_keys_ptr
+    cdef void** c_drv_jit_values_ptr
 
     self._options = options = check_or_create_options(LinkerOptions, options, "Linker options")
 
@@ -496,19 +508,24 @@ cdef inline int Linker_init(Linker self, tuple object_codes, object options) exc
         # the driver writes into via raw pointers during linking operations.
         self._drv_log_bufs = formatted_options
         c_num_opts = len(option_keys)
-        c_jit_keys.resize(c_num_opts)
-        c_jit_values.resize(c_num_opts)
+        self._drv_jit_keys.resize(c_num_opts)
+        self._drv_jit_values.resize(c_num_opts)
         for i in range(c_num_opts):
-            c_jit_keys[i] = <cydriver.CUjit_option><int>option_keys[i]
+            self._drv_jit_keys[i] = <cydriver.CUjit_option><int>option_keys[i]
             val = formatted_options[i]
             if isinstance(val, bytearray):
-                c_jit_values[i] = <void*>PyByteArray_AS_STRING(val)
+                self._drv_jit_values[i] = <void*>PyByteArray_AS_STRING(val)
             else:
-                c_jit_values[i] = <void*><intptr_t>int(val)
+                self._drv_jit_values[i] = <void*><intptr_t>int(val)
+        c_drv_jit_keys_ptr = self._drv_jit_keys.data()
+        c_drv_jit_values_ptr = self._drv_jit_values.data()
         try:
             with nogil:
                 HANDLE_RETURN(cydriver.cuLinkCreate(
-                    <unsigned int>c_num_opts, c_jit_keys.data(), c_jit_values.data(), &c_raw_culink))
+                    <unsigned int>c_num_opts,
+                    c_drv_jit_keys_ptr,
+                    c_drv_jit_values_ptr,
+                    &c_raw_culink))
         except CUDAError as e:
             Linker_annotate_error_log(self, e)
             raise
@@ -622,11 +639,10 @@ cdef inline object Linker_link(Linker self, str target_type):
             raise
         code = (<char*>c_cubin_out)[:c_output_size]
 
-    # Linking is complete; cache the decoded log strings and release
-    # the driver's raw bytearray buffers (no longer written to).
+    # Linking is complete; cache the decoded logs. cuLinkDestroy may still
+    # dereference the raw log-buffer pointers, so retain them until close().
     self._info_log = self.get_info_log()
     self._error_log = self.get_error_log()
-    self._drv_log_bufs = None
 
     return ObjectCode._init(bytes(code), target_type, name=self._options.name)
 
@@ -680,12 +696,22 @@ def _decide_nvjitlink_or_driver() -> bool:
         from cuda.bindings._internal import nvjitlink
 
         if _nvjitlink_has_version_symbol(nvjitlink):
-            _use_nvjitlink_backend = True
-            return False  # Use nvjitlink
-        warn_txt = (
-            f"{'nvJitLink*.dll' if sys.platform == 'win32' else 'libnvJitLink.so*'} is too old (<12.3)."
-            f" Therefore cuda.bindings.nvjitlink is not usable and {warn_txt_common} nvJitLink."
-        )
+            nvjitlink_version = nvjitlink_module.version()
+            driver_major = driver_version()[0]
+            if driver_major <= nvjitlink_version[0]:
+                _use_nvjitlink_backend = True
+                return False  # Use nvjitlink
+
+            warn_txt = (
+                f"CUDA driver major version {driver_major} is newer than "
+                f"nvJitLink major version {nvjitlink_version[0]}; therefore "
+                f"{warn_txt_common} nvJitLink."
+            )
+        else:
+            warn_txt = (
+                f"{'nvJitLink*.dll' if sys.platform == 'win32' else 'libnvJitLink.so*'} is too old (<12.3)."
+                f" Therefore cuda.bindings.nvjitlink is not usable and {warn_txt_common} nvJitLink."
+            )
 
     warn(warn_txt, stacklevel=2, category=RuntimeWarning)
     _use_nvjitlink_backend = False

cuda_core/tests/test_linker.py

@@ -8,6 +8,7 @@
 
 from cuda.core import Device, Linker, LinkerOptions, Program, ProgramOptions, _linker
 from cuda.core._module import ObjectCode
+from cuda.core._program import _can_load_generated_ptx
 from cuda.core._utils.cuda_utils import CUDAError
 
 ARCH = "sm_" + "".join(f"{i}" for i in Device().compute_capability)
@@ -196,19 +197,12 @@ def test_linker_options_as_bytes_nvjitlink():
     assert "-maxrregcount=32" in options_str
 
 
-def test_linker_options_as_bytes_invalid_backend():
+@pytest.mark.parametrize("backend", ("invalid", "driver"))
+def test_linker_options_as_bytes_invalid_backend(backend):
     """Test LinkerOptions.as_bytes() with invalid backend"""
     options = LinkerOptions(arch="sm_80")
     with pytest.raises(ValueError, match="only supports 'nvjitlink' backend"):
-        options.as_bytes("invalid")
-
-
-@pytest.mark.skipif(not is_culink_backend, reason="driver backend test")
-def test_linker_options_as_bytes_driver_not_supported():
-    """Test that as_bytes() is not supported for driver backend"""
-    options = LinkerOptions(arch="sm_80")
-    with pytest.raises(RuntimeError, match="as_bytes\\(\\) only supports 'nvjitlink' backend"):
-        options.as_bytes("driver")
+        options.as_bytes(backend)
 
 
 def test_linker_logs_cached_after_link(compile_ptx_functions):
@@ -234,6 +228,47 @@ def test_linker_handle(compile_ptx_functions):
     assert int(handle) != 0
 
 
+@pytest.mark.agent_authored(model="gpt-5")
+@pytest.mark.skipif(not is_culink_backend, reason="driver backend regression test")
+def test_driver_linker_lifetime_no_heap_corruption(compile_ptx_functions):
+    if not _can_load_generated_ptx():
+        pytest.skip("PTX version too new for current driver")
+
+    linker = Linker(*compile_ptx_functions, options=LinkerOptions(arch=ARCH))
+    linker.link("cubin")
+    linker.close()
+    del linker
+
+    obj_a = Program(kernel_a, "c++", ProgramOptions(relocatable_device_code=True)).compile("ptx")
+    obj_b = Program(device_function_b, "c++", ProgramOptions(relocatable_device_code=True)).compile("ptx")
+    obj_c = Program(device_function_c, "c++", ProgramOptions(relocatable_device_code=True)).compile("ptx")
+    linker = Linker(obj_a, obj_b, obj_c, options=LinkerOptions(arch=ARCH))
+    linker.link("cubin")
+    linker.close()
+
+
+@pytest.mark.agent_authored(model="gpt-5")
+@pytest.mark.skipif(not is_culink_backend, reason="driver backend regression test")
+def test_driver_linker_preserves_error_log_after_close(init_cuda):
+    if not _can_load_generated_ptx():
+        pytest.skip("PTX version too new for current driver")
+
+    bad_kernel = """
+extern __device__ int Z();
+__global__ void A() { int r = Z(); }
+"""
+    bad_obj = Program(bad_kernel, "c++", ProgramOptions(relocatable_device_code=True)).compile("ptx")
+    linker = Linker(bad_obj, options=LinkerOptions(arch=ARCH))
+    with pytest.raises(CUDAError):
+        linker.link("cubin")
+
+    error_log = linker.get_error_log()
+    assert error_log
+    linker.close()
+    assert linker.get_error_log() == error_log
+    assert isinstance(linker.get_info_log(), str)
+
+
 @pytest.mark.skipif(is_culink_backend, reason="nvjitlink options only tested with nvjitlink backend")
 def test_linker_options_nvjitlink_options_as_str():
     """_prepare_nvjitlink_options(as_bytes=False) returns plain strings."""

cuda_core/tests/test_optional_dependency_imports.py

@@ -7,6 +7,14 @@
 from cuda.core import _linker, _program
 
 
+class FakeNvJitLinkModule:
+    def __init__(self, version):
+        self._version = version
+
+    def version(self):
+        return self._version
+
+
 @pytest.fixture(autouse=True)
 def restore_optional_import_state():
     saved_nvvm_module = _program._nvvm_module
@@ -103,3 +111,43 @@ def fake__optional_cuda_import(modname, probe_function=None):
 
     assert use_driver_backend is True
     assert _linker._use_nvjitlink_backend is False
+
+
+@pytest.mark.agent_authored(model="gpt-5")
+@pytest.mark.parametrize(
+    ("driver_version", "nvjitlink_version"),
+    [
+        ((12, 8, 0), (12, 9)),
+        ((12, 9, 0), (13, 0)),
+    ],
+)
+def test_decide_nvjitlink_or_driver_uses_nvjitlink_when_driver_is_not_newer(
+    monkeypatch, driver_version, nvjitlink_version
+):
+    monkeypatch.setattr(
+        _linker,
+        "_optional_cuda_import",
+        lambda *_args, **_kwargs: FakeNvJitLinkModule(nvjitlink_version),
+    )
+    monkeypatch.setattr(_linker, "_nvjitlink_has_version_symbol", lambda _nvjitlink: True)
+    monkeypatch.setattr(_linker, "driver_version", lambda: driver_version)
+
+    assert _linker._decide_nvjitlink_or_driver() is False
+    assert _linker._use_nvjitlink_backend is True
+
+
+@pytest.mark.agent_authored(model="gpt-5")
+def test_decide_nvjitlink_or_driver_falls_back_when_driver_is_newer(monkeypatch):
+    monkeypatch.setattr(
+        _linker,
+        "_optional_cuda_import",
+        lambda *_args, **_kwargs: FakeNvJitLinkModule((12, 9)),
+    )
+    monkeypatch.setattr(_linker, "_nvjitlink_has_version_symbol", lambda _nvjitlink: True)
+    monkeypatch.setattr(_linker, "driver_version", lambda: (13, 0, 0))
+
+    with pytest.warns(RuntimeWarning, match="is newer than nvJitLink major version"):
+        use_driver_backend = _linker._decide_nvjitlink_or_driver()
+
+    assert use_driver_backend is True
+    assert _linker._use_nvjitlink_backend is False

cuda_core/tests/test_program_cache.py

@@ -7,6 +7,10 @@
 
 import pytest
 
+from cuda.core import _linker
+
+is_culink_backend = _linker._decide_nvjitlink_or_driver()
+
 
 def test_program_cache_resource_is_abstract():
     from cuda.core.utils import ProgramCacheResource
@@ -369,9 +373,6 @@ def test_make_program_cache_key_ignores_name_expressions_for_non_nvrtc(code_type
             {"link_time_optimization": None},
             id="lto_false_eq_none",
         ),
-        # ``time`` is a presence gate: the linker emits ``-time`` for any
-        # non-None value, so True / "path" produce the same flag.
-        pytest.param({"time": True}, {"time": "timing.csv"}, id="time_true_eq_path"),
         # ``no_cache`` has an ``is True`` gate; False and None equivalent.
         pytest.param({"no_cache": False}, {"no_cache": None}, id="no_cache_false_eq_none"),
     ],
@@ -390,6 +391,18 @@ def test_make_program_cache_key_ptx_linker_equivalent_options_hash_same(a, b, mo
     assert k_a == k_b
 
 
+@pytest.mark.skipif(is_culink_backend, reason="test requires nvJitLink backend")
+@pytest.mark.agent_authored(model="gpt-5")
+def test_make_program_cache_key_ptx_nvjitlink_time_values_hash_same(monkeypatch):
+    """nvJitLink emits ``-time`` for any non-None value."""
+    from cuda.core.utils import _program_cache
+
+    monkeypatch.setattr(_program_cache._keys, "_linker_backend_and_version", lambda _use_driver: ("nvJitLink", "12030"))
+    k_true = _make_key(code=".version 7.0", code_type="ptx", options=_opts(time=True))
+    k_path = _make_key(code=".version 7.0", code_type="ptx", options=_opts(time="timing.csv"))
+    assert k_true == k_path
+
+
 @pytest.mark.parametrize(
     "field, a, b",
     [
@@ -999,9 +1012,9 @@ def test_make_program_cache_key_rejects_side_effect_options_nvrtc(option_kw, ext
         pytest.param({"time": "whatever.csv"}, id="time_path"),
     ],
 )
+@pytest.mark.skipif(is_culink_backend, reason="test requires nvJitLink backend")
 def test_make_program_cache_key_accepts_side_effect_options_for_ptx(option_kw):
-    """The side-effect guard is NVRTC-specific: PTX (linker) and NVVM must
-    not be blocked by options whose side effects only apply under NVRTC."""
+    """nvJitLink ``-time`` writes only to its info log, not the filesystem."""
     _make_key(code=".version 7.0", code_type="ptx", options=_opts(**option_kw))  # no raise