diff --git a/cuda_core/cuda/core/_linker.pxd b/cuda_core/cuda/core/_linker.pxd index e50ebb9770..1b7d39fd1d 100644 --- a/cuda_core/cuda/core/_linker.pxd +++ b/cuda_core/cuda/core/_linker.pxd @@ -2,6 +2,10 @@ # # SPDX-License-Identifier: Apache-2.0 +from libcpp.vector cimport vector + +from cuda.bindings cimport cydriver + from ._resource_handles cimport NvJitLinkHandle, CuLinkHandle @@ -9,8 +13,12 @@ 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 diff --git a/cuda_core/cuda/core/_linker.pyi b/cuda_core/cuda/core/_linker.pyi index 42b08313f7..37e05fc7fe 100644 --- a/cuda_core/cuda/core/_linker.pyi +++ b/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 ---------- diff --git a/cuda_core/cuda/core/_linker.pyx b/cuda_core/cuda/core/_linker.pyx index 10a85b998a..172b176c84 100644 --- a/cuda_core/cuda/core/_linker.pyx +++ b/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] = option_keys[i] + self._drv_jit_keys[i] = option_keys[i] val = formatted_options[i] if isinstance(val, bytearray): - c_jit_values[i] = PyByteArray_AS_STRING(val) + self._drv_jit_values[i] = PyByteArray_AS_STRING(val) else: - c_jit_values[i] = int(val) + self._drv_jit_values[i] = 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( - c_num_opts, c_jit_keys.data(), c_jit_values.data(), &c_raw_culink)) + 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 = (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 diff --git a/cuda_core/tests/test_linker.py b/cuda_core/tests/test_linker.py index 942ccb1064..9d95b5fd9c 100644 --- a/cuda_core/tests/test_linker.py +++ b/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.""" diff --git a/cuda_core/tests/test_optional_dependency_imports.py b/cuda_core/tests/test_optional_dependency_imports.py index 02edcc9839..6759fd3cfa 100644 --- a/cuda_core/tests/test_optional_dependency_imports.py +++ b/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 diff --git a/cuda_core/tests/test_program_cache.py b/cuda_core/tests/test_program_cache.py index 963ec1cc04..3bb579e2fe 100644 --- a/cuda_core/tests/test_program_cache.py +++ b/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