_linker.pyx

# SPDX-FileCopyrightText: Copyright (c) 2024-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
#
# SPDX-License-Identifier: Apache-2.0
"""Linking machinery for combining object codes.

This module provides :class:`Linker` for linking one or more
:class:`~cuda.core.ObjectCode` objects, with :class:`LinkerOptions` for
configuration.
"""

from __future__ import annotations

from cpython.bytearray cimport PyByteArray_AS_STRING
from libc.stdint cimport intptr_t, uint32_t
from libcpp.vector cimport vector
from cuda.bindings cimport cydriver
from cuda.bindings cimport cynvjitlink

from ._resource_handles cimport (
    as_cu,
    as_py,
    create_culink_handle,
    create_nvjitlink_handle,
)
from cuda.core._utils.cuda_utils cimport HANDLE_RETURN, HANDLE_RETURN_NVJITLINK

import sys
from dataclasses import dataclass
from typing import Union
from warnings import warn

from cuda.pathfinder._optional_cuda_import import _optional_cuda_import
from cuda.core._device import Device
from cuda.core._module import ObjectCode
from cuda.core._utils.clear_error_support import assert_type
from cuda.core._utils.cuda_utils import (
    CUDAError,
    check_or_create_options,
    driver,
    is_sequence,
)

ctypedef const char* const_char_ptr
ctypedef void* void_ptr

__all__ = ["Linker", "LinkerOptions"]

LinkerHandleT = Union["cuda.bindings.nvjitlink.nvJitLinkHandle", "cuda.bindings.driver.CUlinkState"]


# =============================================================================
# Principal class
# =============================================================================

cdef class Linker:
    """Represent a linking machinery to link one or more object codes into
    :class:`~cuda.core.ObjectCode`.

    This object provides a unified interface to multiple underlying
    linker libraries (such as nvJitLink or cuLink* from the CUDA driver).

    Parameters
    ----------
    object_codes : :class:`~cuda.core.ObjectCode`
        One or more ObjectCode objects to be linked.
    options : :class:`LinkerOptions`, optional
        Options for the linker. If not provided, default options will be used.
    """

    def __init__(self, *object_codes: ObjectCode, options: "LinkerOptions" = None):
        Linker_init(self, object_codes, options)

    def link(self, target_type) -> ObjectCode:
        """Link the provided object codes into a single output of the specified target type.

        Parameters
        ----------
        target_type : str
            The type of the target output. Must be either "cubin" or "ptx".

        Returns
        -------
        :class:`~cuda.core.ObjectCode`
            The linked object code of the specified target type.

        .. note::

            Ensure that input object codes were compiled with appropriate
            flags for linking (e.g., relocatable device code enabled).
        """
        return Linker_link(self, target_type)

    def get_error_log(self) -> str:
        """Get the error log generated by the linker.

        Returns
        -------
        str
            The error log.
        """
        # After link(), the decoded log is cached here.
        if self._error_log is not None:
            return self._error_log
        cdef cynvjitlink.nvJitLinkHandle c_h
        cdef size_t c_log_size = 0
        cdef char* c_log_ptr
        if self._use_nvjitlink:
            c_h = as_cu(self._nvjitlink_handle)
            cynvjitlink.nvJitLinkGetErrorLogSize(c_h, &c_log_size)
            log = bytearray(c_log_size)
            if c_log_size > 0:
                c_log_ptr = <char*>(<bytearray>log)
                cynvjitlink.nvJitLinkGetErrorLog(c_h, c_log_ptr)
            return log.decode("utf-8", errors="backslashreplace")
        else:
            return (<bytearray>self._drv_log_bufs[2]).decode(
                "utf-8", errors="backslashreplace").rstrip('\x00')

    def get_info_log(self) -> str:
        """Get the info log generated by the linker.

        Returns
        -------
        str
            The info log.
        """
        # After link(), the decoded log is cached here.
        if self._info_log is not None:
            return self._info_log
        cdef cynvjitlink.nvJitLinkHandle c_h
        cdef size_t c_log_size = 0
        cdef char* c_log_ptr
        if self._use_nvjitlink:
            c_h = as_cu(self._nvjitlink_handle)
            cynvjitlink.nvJitLinkGetInfoLogSize(c_h, &c_log_size)
            log = bytearray(c_log_size)
            if c_log_size > 0:
                c_log_ptr = <char*>(<bytearray>log)
                cynvjitlink.nvJitLinkGetInfoLog(c_h, c_log_ptr)
            return log.decode("utf-8", errors="backslashreplace")
        else:
            return (<bytearray>self._drv_log_bufs[0]).decode(
                "utf-8", errors="backslashreplace").rstrip('\x00')

    def close(self):
        """Destroy this linker."""
        if self._use_nvjitlink:
            self._nvjitlink_handle.reset()
        else:
            self._culink_handle.reset()

    @property
    def handle(self) -> LinkerHandleT:
        """Return the underlying handle object.

        .. note::

           The type of the returned object depends on the backend.

        .. caution::

            This handle is a Python object. To get the memory address of the underlying C
            handle, call ``int(Linker.handle)``.
        """
        if self._use_nvjitlink:
            return as_py(self._nvjitlink_handle)
        else:
            return as_py(self._culink_handle)

    @classmethod
    def backend(cls) -> str:
        """Return which linking backend will be used.

        Returns ``"nvJitLink"`` when the nvJitLink library is available
        and meets the minimum version requirement, otherwise ``"driver"``.

        .. note::

            Prefer letting :class:`Linker` decide. Query ``backend()``
            only when you need to dispatch based on input format (for
            example: choose PTX vs. LTOIR before constructing a
            ``Linker``). The returned string names an implementation
            detail whose support matrix may shift across CTK releases.
        """
        return "driver" if _decide_nvjitlink_or_driver() else "nvJitLink"


# =============================================================================
# Supporting classes
# =============================================================================

@dataclass
class LinkerOptions:
    """Customizable options for configuring :class:`Linker`.

    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.

    Attributes
    ----------
    name : str, optional
        Name of the linker. If the linking succeeds, the name is passed down to the generated `ObjectCode`.
    arch : str, optional
        Pass the SM architecture value, such as ``sm_<CC>`` (for generating CUBIN) or
        ``compute_<CC>`` (for generating PTX). If not provided, the current device's architecture
        will be used.
    max_register_count : int, optional
        Maximum register count.
    time : bool, optional
        Print timing information to the info log.
        Default: False.
    verbose : bool, optional
        Print verbose messages to the info log.
        Default: False.
    link_time_optimization : bool, optional
        Perform link time optimization.
        Default: False.
    ptx : bool, optional
        Emit PTX after linking instead of CUBIN; only supported with ``link_time_optimization=True``.
        Default: False.
    optimization_level : int, optional
        Set optimization level. Only 0 and 3 are accepted.
    debug : bool, optional
        Generate debug information.
        Default: False.
    lineinfo : bool, optional
        Generate line information.
        Default: False.
    ftz : bool, optional
        Flush denormal values to zero.
        Default: False.
    prec_div : bool, optional
        Use precise division.
        Default: True.
    prec_sqrt : bool, optional
        Use precise square root.
        Default: True.
    fma : bool, optional
        Use fast multiply-add.
        Default: True.
    kernels_used : [str | tuple[str] | list[str]], optional
        Pass a kernel or sequence of kernels that are used; any not in the list can be removed.
    variables_used : [str | tuple[str] | list[str]], optional
        Pass a variable or sequence of variables that are used; any not in the list can be removed.
    optimize_unused_variables : bool, optional
        Assume that if a variable is not referenced in device code, it can be removed.
        Default: False.
    ptxas_options : [str | tuple[str] | list[str]], optional
        Pass options to PTXAS.
    split_compile : int, optional
        Split compilation maximum thread count. Use 0 to use all available processors. Value of 1 disables split
        compilation (default).
        Default: 1.
    split_compile_extended : int, optional
        A more aggressive form of split compilation available in LTO mode only. Accepts a maximum thread count value.
        Use 0 to use all available processors. Value of 1 disables extended split compilation (default). Note: This
        option can potentially impact performance of the compiled binary.
        Default: 1.
    no_cache : bool, optional
        Do not cache the intermediate steps of nvJitLink.
        Default: False.
    """

    name: str | None = "<default linker>"
    arch: str | None = None
    max_register_count: int | None = None
    time: bool | None = None
    verbose: bool | None = None
    link_time_optimization: bool | None = None
    ptx: bool | None = None
    optimization_level: int | None = None
    debug: bool | None = None
    lineinfo: bool | None = None
    ftz: bool | None = None
    prec_div: bool | None = None
    prec_sqrt: bool | None = None
    fma: bool | None = None
    kernels_used: str | tuple[str] | list[str] | None = None
    variables_used: str | tuple[str] | list[str] | None = None
    optimize_unused_variables: bool | None = None
    ptxas_options: str | tuple[str] | list[str] | None = None
    split_compile: int | None = None
    split_compile_extended: int | None = None
    no_cache: bool | None = None

    def __post_init__(self):
        _lazy_init()
        self._name = self.name.encode()

    def _prepare_nvjitlink_options(self, as_bytes: bool = False) -> list[bytes] | list[str]:
        options = []

        if self.arch is not None:
            options.append(f"-arch={self.arch}")
        else:
            options.append("-arch=sm_" + "".join(f"{i}" for i in Device().compute_capability))
        if self.max_register_count is not None:
            options.append(f"-maxrregcount={self.max_register_count}")
        if self.time is not None:
            options.append("-time")
        if self.verbose:
            options.append("-verbose")
        if self.link_time_optimization:
            options.append("-lto")
        if self.ptx:
            options.append("-ptx")
        if self.optimization_level is not None:
            options.append(f"-O{self.optimization_level}")
        if self.debug:
            options.append("-g")
        if self.lineinfo:
            options.append("-lineinfo")
        if self.ftz is not None:
            options.append(f"-ftz={'true' if self.ftz else 'false'}")
        if self.prec_div is not None:
            options.append(f"-prec-div={'true' if self.prec_div else 'false'}")
        if self.prec_sqrt is not None:
            options.append(f"-prec-sqrt={'true' if self.prec_sqrt else 'false'}")
        if self.fma is not None:
            options.append(f"-fma={'true' if self.fma else 'false'}")
        if self.kernels_used is not None:
            if isinstance(self.kernels_used, str):
                options.append(f"-kernels-used={self.kernels_used}")
            elif isinstance(self.kernels_used, list):
                for kernel in self.kernels_used:
                    options.append(f"-kernels-used={kernel}")
        if self.variables_used is not None:
            if isinstance(self.variables_used, str):
                options.append(f"-variables-used={self.variables_used}")
            elif isinstance(self.variables_used, list):
                for variable in self.variables_used:
                    options.append(f"-variables-used={variable}")
        if self.optimize_unused_variables is not None:
            options.append("-optimize-unused-variables")
        if self.ptxas_options is not None:
            if isinstance(self.ptxas_options, str):
                options.append(f"-Xptxas={self.ptxas_options}")
            elif is_sequence(self.ptxas_options):
                for opt in self.ptxas_options:
                    options.append(f"-Xptxas={opt}")
        if self.split_compile is not None:
            options.append(f"-split-compile={self.split_compile}")
        if self.split_compile_extended is not None:
            options.append(f"-split-compile-extended={self.split_compile_extended}")
        if self.no_cache is True:
            options.append("-no-cache")

        if as_bytes:
            return [o.encode() for o in options]
        else:
            return options

    def _prepare_driver_options(self) -> tuple[list, list]:
        formatted_options = []
        option_keys = []

        # allocate a fixed-sized buffer for each info/error log
        size = 4194304
        formatted_options.extend((bytearray(size), size, bytearray(size), size))
        option_keys.extend(
            (
                _driver.CUjit_option.CU_JIT_INFO_LOG_BUFFER,
                _driver.CUjit_option.CU_JIT_INFO_LOG_BUFFER_SIZE_BYTES,
                _driver.CUjit_option.CU_JIT_ERROR_LOG_BUFFER,
                _driver.CUjit_option.CU_JIT_ERROR_LOG_BUFFER_SIZE_BYTES,
            )
        )

        if self.arch is not None:
            arch = self.arch.split("_")[-1].upper()
            formatted_options.append(getattr(_driver.CUjit_target, f"CU_TARGET_COMPUTE_{arch}"))
            option_keys.append(_driver.CUjit_option.CU_JIT_TARGET)
        if self.max_register_count is not None:
            formatted_options.append(self.max_register_count)
            option_keys.append(_driver.CUjit_option.CU_JIT_MAX_REGISTERS)
        if self.time is not None:
            raise ValueError("time option is not supported by the driver API")
        if self.verbose:
            formatted_options.append(1)
            option_keys.append(_driver.CUjit_option.CU_JIT_LOG_VERBOSE)
        if self.link_time_optimization:
            formatted_options.append(1)
            option_keys.append(_driver.CUjit_option.CU_JIT_LTO)
        if self.ptx:
            raise ValueError("ptx option is not supported by the driver API")
        if self.optimization_level is not None:
            formatted_options.append(self.optimization_level)
            option_keys.append(_driver.CUjit_option.CU_JIT_OPTIMIZATION_LEVEL)
        if self.debug:
            formatted_options.append(1)
            option_keys.append(_driver.CUjit_option.CU_JIT_GENERATE_DEBUG_INFO)
        if self.lineinfo:
            formatted_options.append(1)
            option_keys.append(_driver.CUjit_option.CU_JIT_GENERATE_LINE_INFO)
        if self.ftz is not None:
            warn("ftz option is deprecated in the driver API", DeprecationWarning, stacklevel=3)
        if self.prec_div is not None:
            warn("prec_div option is deprecated in the driver API", DeprecationWarning, stacklevel=3)
        if self.prec_sqrt is not None:
            warn("prec_sqrt option is deprecated in the driver API", DeprecationWarning, stacklevel=3)
        if self.fma is not None:
            warn("fma options is deprecated in the driver API", DeprecationWarning, stacklevel=3)
        if self.kernels_used is not None:
            warn("kernels_used is deprecated in the driver API", DeprecationWarning, stacklevel=3)
        if self.variables_used is not None:
            warn("variables_used is deprecated in the driver API", DeprecationWarning, stacklevel=3)
        if self.optimize_unused_variables is not None:
            warn("optimize_unused_variables is deprecated in the driver API", DeprecationWarning, stacklevel=3)
        if self.ptxas_options is not None:
            raise ValueError("ptxas_options option is not supported by the driver API")
        if self.split_compile is not None:
            raise ValueError("split_compile option is not supported by the driver API")
        if self.split_compile_extended is not None:
            raise ValueError("split_compile_extended option is not supported by the driver API")
        if self.no_cache is True:
            formatted_options.append(_driver.CUjit_cacheMode.CU_JIT_CACHE_OPTION_NONE)
            option_keys.append(_driver.CUjit_option.CU_JIT_CACHE_MODE)

        return formatted_options, option_keys

    def as_bytes(self, backend: str = "nvjitlink") -> list[bytes]:
        """Convert linker options to bytes format for the nvjitlink backend.

        Parameters
        ----------
        backend : str, optional
            The linker backend. Only "nvjitlink" is supported. Default is "nvjitlink".

        Returns
        -------
        list[bytes]
            List of option strings encoded as bytes.

        Raises
        ------
        ValueError
            If an unsupported backend is specified.
        RuntimeError
            If nvJitLink backend is not available.
        """
        backend = backend.lower()
        if backend != "nvjitlink":
            raise ValueError(f"as_bytes() only supports 'nvjitlink' backend, got '{backend}'")
        if not _use_nvjitlink_backend:
            raise RuntimeError("nvJitLink backend is not available")
        return self._prepare_nvjitlink_options(as_bytes=True)


# =============================================================================
# Private implementation: cdef inline helpers
# =============================================================================

cdef inline int Linker_init(Linker self, tuple object_codes, object options) except -1:
    """Initialize a Linker instance."""
    if len(object_codes) == 0:
        raise ValueError("At least one ObjectCode object must be provided")

    cdef cynvjitlink.nvJitLinkHandle c_raw_nvjitlink
    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

    self._options = options = check_or_create_options(LinkerOptions, options, "Linker options")

    if _use_nvjitlink_backend:
        self._use_nvjitlink = True
        options_bytes = options._prepare_nvjitlink_options(as_bytes=True)
        c_num_opts = len(options_bytes)
        c_str_opts.resize(c_num_opts)
        for i in range(c_num_opts):
            c_str_opts[i] = <const char*>(<bytes>options_bytes[i])
        with nogil:
            HANDLE_RETURN_NVJITLINK(NULL, cynvjitlink.nvJitLinkCreate(
                &c_raw_nvjitlink, <uint32_t>c_num_opts, c_str_opts.data()))
        self._nvjitlink_handle = create_nvjitlink_handle(c_raw_nvjitlink)
    else:
        self._use_nvjitlink = False
        formatted_options, option_keys = options._prepare_driver_options()
        # Keep the formatted_options list alive: it contains bytearrays that
        # 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)
        for i in range(c_num_opts):
            c_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)
            else:
                c_jit_values[i] = <void*><intptr_t>int(val)
        try:
            with nogil:
                HANDLE_RETURN(cydriver.cuLinkCreate(
                    <unsigned int>c_num_opts, c_jit_keys.data(), c_jit_values.data(), &c_raw_culink))
        except CUDAError as e:
            Linker_annotate_error_log(self, e)
            raise
        self._culink_handle = create_culink_handle(c_raw_culink)

    for code in object_codes:
        assert_type(code, ObjectCode)
        Linker_add_code_object(self, code)
    return 0


cdef inline void Linker_add_code_object(Linker self, object object_code) except *:
    """Add a single ObjectCode to the linker."""
    data = object_code.code
    cdef cynvjitlink.nvJitLinkHandle c_nvjitlink_h
    cdef cydriver.CUlinkState c_culink_state
    cdef cynvjitlink.nvJitLinkInputType c_nv_input_type
    cdef cydriver.CUjitInputType c_drv_input_type
    cdef const char* c_data_ptr
    cdef size_t c_data_size
    cdef const char* c_name_ptr
    cdef const char* c_file_ptr

    name_bytes = f"{object_code.name}".encode()
    c_name_ptr = <const char*>name_bytes

    input_types = _nvjitlink_input_types if self._use_nvjitlink else _driver_input_types
    py_input_type = input_types.get(object_code.code_type)
    if py_input_type is None:
        raise ValueError(f"Unknown code_type associated with ObjectCode: {object_code.code_type}")

    if self._use_nvjitlink:
        c_nvjitlink_h = as_cu(self._nvjitlink_handle)
        c_nv_input_type = <cynvjitlink.nvJitLinkInputType><int>py_input_type
        if isinstance(data, bytes):
            c_data_ptr = <const char*>(<bytes>data)
            c_data_size = len(data)
            with nogil:
                HANDLE_RETURN_NVJITLINK(c_nvjitlink_h, cynvjitlink.nvJitLinkAddData(
                    c_nvjitlink_h, c_nv_input_type, <const void*>c_data_ptr, c_data_size, c_name_ptr))
        elif isinstance(data, str):
            file_bytes = data.encode()
            c_file_ptr = <const char*>file_bytes
            with nogil:
                HANDLE_RETURN_NVJITLINK(c_nvjitlink_h, cynvjitlink.nvJitLinkAddFile(
                    c_nvjitlink_h, c_nv_input_type, c_file_ptr))
        else:
            raise TypeError(f"Expected bytes or str, but got {type(data).__name__}")
    else:
        c_culink_state = as_cu(self._culink_handle)
        c_drv_input_type = <cydriver.CUjitInputType><int>py_input_type
        try:
            if isinstance(data, bytes):
                c_data_ptr = <const char*>(<bytes>data)
                c_data_size = len(data)
                with nogil:
                    HANDLE_RETURN(cydriver.cuLinkAddData(
                        c_culink_state, c_drv_input_type, <void*>c_data_ptr, c_data_size, c_name_ptr,
                        0, NULL, NULL))
            elif isinstance(data, str):
                file_bytes = data.encode()
                c_file_ptr = <const char*>file_bytes
                with nogil:
                    HANDLE_RETURN(cydriver.cuLinkAddFile(
                        c_culink_state, c_drv_input_type, c_file_ptr, 0, NULL, NULL))
            else:
                raise TypeError(f"Expected bytes or str, but got {type(data).__name__}")
        except CUDAError as e:
            Linker_annotate_error_log(self, e)
            raise


cdef inline object Linker_link(Linker self, str target_type):
    """Complete linking and return the result as ObjectCode."""
    if target_type not in ("cubin", "ptx"):
        raise ValueError(f"Unsupported target type: {target_type}")

    cdef cynvjitlink.nvJitLinkHandle c_nvjitlink_h
    cdef cydriver.CUlinkState c_culink_state
    cdef size_t c_output_size = 0
    cdef char* c_code_ptr
    cdef void* c_cubin_out = NULL

    if self._use_nvjitlink:
        c_nvjitlink_h = as_cu(self._nvjitlink_handle)
        with nogil:
            HANDLE_RETURN_NVJITLINK(c_nvjitlink_h, cynvjitlink.nvJitLinkComplete(c_nvjitlink_h))
        if target_type == "cubin":
            HANDLE_RETURN_NVJITLINK(c_nvjitlink_h,
                cynvjitlink.nvJitLinkGetLinkedCubinSize(c_nvjitlink_h, &c_output_size))
            code = bytearray(c_output_size)
            c_code_ptr = <char*>(<bytearray>code)
            with nogil:
                HANDLE_RETURN_NVJITLINK(c_nvjitlink_h,
                    cynvjitlink.nvJitLinkGetLinkedCubin(c_nvjitlink_h, c_code_ptr))
        else:
            HANDLE_RETURN_NVJITLINK(c_nvjitlink_h,
                cynvjitlink.nvJitLinkGetLinkedPtxSize(c_nvjitlink_h, &c_output_size))
            code = bytearray(c_output_size)
            c_code_ptr = <char*>(<bytearray>code)
            with nogil:
                HANDLE_RETURN_NVJITLINK(c_nvjitlink_h,
                    cynvjitlink.nvJitLinkGetLinkedPtx(c_nvjitlink_h, c_code_ptr))
    else:
        c_culink_state = as_cu(self._culink_handle)
        try:
            with nogil:
                HANDLE_RETURN(cydriver.cuLinkComplete(c_culink_state, &c_cubin_out, &c_output_size))
        except CUDAError as e:
            Linker_annotate_error_log(self, e)
            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).
    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)


cdef inline void Linker_annotate_error_log(Linker self, object e):
    """Annotate a CUDAError with the driver linker error log."""
    error_log = self.get_error_log()
    if error_log:
        e.args = (e.args[0] + f"\nLinker error log: {error_log}", *e.args[1:])


# =============================================================================
# Private implementation: module-level state and initialization
# =============================================================================

# TODO: revisit this treatment for py313t builds
_driver = None  # populated if nvJitLink cannot be used
_inited = False
_use_nvjitlink_backend = None  # set by _decide_nvjitlink_or_driver()

# Input type mappings populated by _lazy_init() with C-level enum ints.
_nvjitlink_input_types = None
_driver_input_types = None


def _nvjitlink_has_version_symbol(nvjitlink) -> bool:
    # This condition is equivalent to testing for version >= 12.3
    return bool(nvjitlink._inspect_function_pointer("__nvJitLinkVersion"))


# Note: this function is reused in the tests
def _decide_nvjitlink_or_driver() -> bool:
    """Return True if falling back to the cuLink* driver APIs."""
    global _driver, _use_nvjitlink_backend
    if _use_nvjitlink_backend is not None:
        return not _use_nvjitlink_backend

    warn_txt_common = (
        "the driver APIs will be used instead, which do not support"
        " minor version compatibility or linking LTO IRs."
        " For best results, consider upgrading to a recent version of"
    )

    nvjitlink_module = _optional_cuda_import(
        "cuda.bindings.nvjitlink",
        probe_function=lambda module: module.version(),  # probe triggers nvJitLink runtime load
    )
    if nvjitlink_module is None:
        warn_txt = f"cuda.bindings.nvjitlink is not available, therefore {warn_txt_common} cuda-bindings."
    else:
        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."
        )

    warn(warn_txt, stacklevel=2, category=RuntimeWarning)
    _use_nvjitlink_backend = False
    _driver = driver
    return True


def _lazy_init():
    global _inited, _nvjitlink_input_types, _driver_input_types
    if _inited:
        return

    _decide_nvjitlink_or_driver()
    if _use_nvjitlink_backend:
        _nvjitlink_input_types = {
            "ptx": <int>cynvjitlink.NVJITLINK_INPUT_PTX,
            "cubin": <int>cynvjitlink.NVJITLINK_INPUT_CUBIN,
            "fatbin": <int>cynvjitlink.NVJITLINK_INPUT_FATBIN,
            "ltoir": <int>cynvjitlink.NVJITLINK_INPUT_LTOIR,
            "object": <int>cynvjitlink.NVJITLINK_INPUT_OBJECT,
            "library": <int>cynvjitlink.NVJITLINK_INPUT_LIBRARY,
        }
    else:
        _driver_input_types = {
            "ptx": <int>cydriver.CU_JIT_INPUT_PTX,
            "cubin": <int>cydriver.CU_JIT_INPUT_CUBIN,
            "fatbin": <int>cydriver.CU_JIT_INPUT_FATBINARY,
            "object": <int>cydriver.CU_JIT_INPUT_OBJECT,
            "library": <int>cydriver.CU_JIT_INPUT_LIBRARY,
        }
    _inited = True