test_linker.py

# SPDX-FileCopyrightText: Copyright (c) 2024-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
#
# SPDX-License-Identifier: Apache-2.0

import inspect

import pytest

from cuda.core import Device, Linker, LinkerOptions, Program, ProgramOptions, _linker
from cuda.core._module import ObjectCode
from cuda.core._utils.cuda_utils import CUDAError

ARCH = "sm_" + "".join(f"{i}" for i in Device().compute_capability)

kernel_a = """
extern __device__ int B();
extern __device__ int C(int a, int b);
__global__ void A() { int result = C(B(), 1);}
"""
device_function_b = "__device__ int B() { return 0; }"
device_function_c = "__device__ int C(int a, int b) { return a + b; }"

is_culink_backend = _linker._decide_nvjitlink_or_driver()
if not is_culink_backend:
    from cuda.bindings import nvjitlink

    nvJitLinkError = nvjitlink.nvJitLinkError
else:

    class nvJitLinkError(Exception):
        pass


@pytest.fixture
def compile_ptx_functions(init_cuda):
    # Without -rdc (relocatable device code) option, the generated ptx will not included any unreferenced
    # device functions, causing the link to fail
    object_code_a_ptx = Program(kernel_a, "c++", ProgramOptions(relocatable_device_code=True)).compile("ptx")
    object_code_b_ptx = Program(device_function_b, "c++", ProgramOptions(relocatable_device_code=True)).compile("ptx")
    object_code_c_ptx = Program(device_function_c, "c++", ProgramOptions(relocatable_device_code=True)).compile("ptx")

    return object_code_a_ptx, object_code_b_ptx, object_code_c_ptx


@pytest.fixture
def compile_ltoir_functions(init_cuda):
    object_code_a_ltoir = Program(kernel_a, "c++", ProgramOptions(link_time_optimization=True)).compile("ltoir")
    object_code_b_ltoir = Program(device_function_b, "c++", ProgramOptions(link_time_optimization=True)).compile(
        "ltoir"
    )
    object_code_c_ltoir = Program(device_function_c, "c++", ProgramOptions(link_time_optimization=True)).compile(
        "ltoir"
    )

    return object_code_a_ltoir, object_code_b_ltoir, object_code_c_ltoir


options = [
    LinkerOptions(),
    LinkerOptions(arch=ARCH, verbose=True),
    LinkerOptions(arch=ARCH, max_register_count=32),
    LinkerOptions(arch=ARCH, optimization_level=3),
    LinkerOptions(arch=ARCH, debug=True),
    LinkerOptions(arch=ARCH, lineinfo=True),
]
if not is_culink_backend:
    options += [
        LinkerOptions(arch=ARCH, time=True),
        LinkerOptions(arch=ARCH, optimize_unused_variables=True),
        LinkerOptions(arch=ARCH, ptxas_options="-v"),
        LinkerOptions(arch=ARCH, ptxas_options=["-v", "--verbose"]),
        LinkerOptions(arch=ARCH, ptxas_options=("-v", "--verbose")),
        LinkerOptions(arch=ARCH, split_compile=0),
        LinkerOptions(arch=ARCH, split_compile_extended=1),
        # The following options are supported by nvjitlink and deprecated by culink
        LinkerOptions(arch=ARCH, ftz=True),
        LinkerOptions(arch=ARCH, prec_div=True),
        LinkerOptions(arch=ARCH, prec_sqrt=True),
        LinkerOptions(arch=ARCH, fma=True),
        LinkerOptions(arch=ARCH, kernels_used="A"),
        LinkerOptions(arch=ARCH, kernels_used=["C", "B"]),
        LinkerOptions(arch=ARCH, kernels_used=("C", "B")),
        LinkerOptions(arch=ARCH, variables_used="var1"),
        LinkerOptions(arch=ARCH, variables_used=["var1", "var2"]),
        LinkerOptions(arch=ARCH, variables_used=("var1", "var2")),
    ]
    version = nvjitlink.version()
    if version >= (12, 5):
        options.append(LinkerOptions(arch=ARCH, no_cache=True))


@pytest.mark.parametrize("options", options)
def test_linker_init(compile_ptx_functions, options):
    linker = Linker(*compile_ptx_functions, options=options)
    object_code = linker.link("cubin")
    assert isinstance(object_code, ObjectCode)
    assert linker.backend() == ("driver" if is_culink_backend else "nvJitLink")


def test_linker_init_invalid_arch(compile_ptx_functions):
    err = AttributeError if is_culink_backend else nvjitlink.nvJitLinkError
    with pytest.raises(err):
        options = LinkerOptions(arch="99", ptx=True)
        Linker(*compile_ptx_functions, options=options)


@pytest.mark.skipif(is_culink_backend, reason="culink does not support ptx option")
def test_linker_link_ptx_nvjitlink(compile_ltoir_functions):
    options = LinkerOptions(arch=ARCH, link_time_optimization=True, ptx=True)
    linker = Linker(*compile_ltoir_functions, options=options)
    linked_code = linker.link("ptx")
    assert isinstance(linked_code, ObjectCode)
    assert linked_code.name == options.name


@pytest.mark.skipif(not is_culink_backend, reason="nvjitlink requires lto for ptx linking")
def test_linker_link_ptx_culink(compile_ptx_functions):
    options = LinkerOptions(arch=ARCH)
    linker = Linker(*compile_ptx_functions, options=options)
    linked_code = linker.link("ptx")
    assert isinstance(linked_code, ObjectCode)
    assert linked_code.name == options.name


def test_linker_link_cubin(compile_ptx_functions):
    options = LinkerOptions(arch=ARCH)
    linker = Linker(*compile_ptx_functions, options=options)
    linked_code = linker.link("cubin")
    assert isinstance(linked_code, ObjectCode)
    assert linked_code.name == options.name


def test_linker_link_ptx_multiple(compile_ptx_functions):
    ptxes = tuple(ObjectCode.from_ptx(obj.code) for obj in compile_ptx_functions)
    options = LinkerOptions(arch=ARCH)
    linker = Linker(*ptxes, options=options)
    linked_code = linker.link("cubin")
    assert isinstance(linked_code, ObjectCode)
    assert linked_code.name == options.name


def test_linker_link_invalid_target_type(compile_ptx_functions):
    options = LinkerOptions(arch=ARCH)
    linker = Linker(*compile_ptx_functions, options=options)
    with pytest.raises(ValueError):
        linker.link("invalid_target")


def test_linker_get_error_log(compile_ptx_functions):
    options = LinkerOptions(name="ABC", arch=ARCH)

    replacement_kernel = """
extern __device__ int Z();
extern __device__ int C(int a, int b);
__global__ void A() { int result = C(Z(), 1);}
"""
    dummy_program = Program(
        replacement_kernel, "c++", ProgramOptions(name="CBA", relocatable_device_code=True)
    ).compile("ptx")
    linker = Linker(dummy_program, *(compile_ptx_functions[1:]), options=options)
    try:
        linker.link("cubin")

    except (nvJitLinkError, CUDAError):
        log = linker.get_error_log()
        assert isinstance(log, str)
        # TODO when 4902246 is addressed, we can update this to cover nvjitlink as well
        # The error is coming from the input object that's being linked (CBA), not the output object (ABC).
        if is_culink_backend:
            assert log.rstrip("\x00") == "error   : Undefined reference to '_Z1Zv' in 'CBA'"


def test_linker_get_info_log(compile_ptx_functions):
    options = LinkerOptions(arch=ARCH)
    linker = Linker(*compile_ptx_functions, options=options)
    linker.link("cubin")
    log = linker.get_info_log()
    assert isinstance(log, str)


@pytest.mark.skipif(is_culink_backend, reason="as_bytes() only supported for nvjitlink backend")
def test_linker_options_as_bytes_nvjitlink():
    """Test LinkerOptions.as_bytes() for nvJitLink backend"""
    options = LinkerOptions(arch="sm_80", debug=True, ftz=True, max_register_count=32)
    nvjitlink_options = options.as_bytes("nvjitlink")

    # Should return list of bytes
    assert isinstance(nvjitlink_options, list)
    assert all(isinstance(opt, bytes) for opt in nvjitlink_options)

    # Decode to check content
    options_str = [opt.decode() for opt in nvjitlink_options]
    assert "-arch=sm_80" in options_str
    assert "-g" in options_str
    assert "-ftz=true" in options_str
    assert "-maxrregcount=32" in options_str


def test_linker_options_as_bytes_invalid_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")


def test_linker_logs_cached_after_link(compile_ptx_functions):
    """After a successful link(), get_error_log/get_info_log should return cached strings."""
    options = LinkerOptions(arch=ARCH)
    linker = Linker(*compile_ptx_functions, options=options)
    linker.link("cubin")
    err_log = linker.get_error_log()
    info_log = linker.get_info_log()
    assert isinstance(err_log, str)
    assert isinstance(info_log, str)
    # Calling again should return the same observable values.
    assert linker.get_error_log() == err_log
    assert linker.get_info_log() == info_log


def test_linker_handle(compile_ptx_functions):
    """Linker.handle returns a non-null handle object."""
    options = LinkerOptions(arch=ARCH)
    linker = Linker(*compile_ptx_functions, options=options)
    handle = linker.handle
    assert handle is not None
    assert int(handle) != 0


@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."""
    opts = LinkerOptions(arch=ARCH, debug=True, lineinfo=True)
    options = opts._prepare_nvjitlink_options(as_bytes=False)
    assert isinstance(options, list)
    assert all(isinstance(o, str) for o in options)
    assert f"-arch={ARCH}" in options
    assert "-g" in options
    assert "-lineinfo" in options


class TestBackendClassmethod:
    def test_backend_returns_nvjitlink(self, monkeypatch):
        monkeypatch.setattr(_linker, "_use_nvjitlink_backend", True)
        assert Linker.backend() == "nvJitLink"

    def test_backend_returns_driver(self, monkeypatch):
        monkeypatch.setattr(_linker, "_use_nvjitlink_backend", False)
        assert Linker.backend() == "driver"

    def test_backend_invokes_probe_when_not_memoised(self, monkeypatch):
        monkeypatch.setattr(_linker, "_use_nvjitlink_backend", None)
        called = []

        def fake_decide():
            called.append(True)
            return False  # False = not falling back to driver = nvJitLink

        monkeypatch.setattr(_linker, "_decide_nvjitlink_or_driver", fake_decide)
        result = Linker.backend()
        assert result == "nvJitLink"
        assert called, "_decide_nvjitlink_or_driver was not called"

    def test_backend_is_classmethod(self):
        attr = inspect.getattr_static(Linker, "backend")
        assert isinstance(attr, classmethod)

    def test_backend_is_not_property(self):
        """backend is a classmethod, not a property.

        This is an intentional breaking change from the prior property API.
        Attribute-style access (``linker.backend``) now returns a bound method,
        not a string. All call sites must use parens: ``Linker.backend()``.
        """
        attr = inspect.getattr_static(Linker, "backend")
        assert not isinstance(attr, property)