feat(cuda.core): support nvJitLink incremental linking (#2867)

Author: isVoid

cuda_core/cuda/core/_linker.pxd

@@ -18,6 +18,7 @@ cdef class Linker:
         vector[cydriver.CUjit_option] _drv_jit_keys
         vector[void*] _drv_jit_values
         bint _use_nvjitlink
+        bint _has_ptx_or_cubin_input
         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()

cuda_core/cuda/core/_linker.pyi

@@ -21,6 +21,8 @@ const_char_ptr: TypeAlias = bytes
 __all__ = ['Linker', 'LinkerOptions']
 LinkerHandleT = Union['cuda.bindings.nvjitlink.nvJitLinkHandle', 'cuda.bindings.driver.CUlinkState']
 _driver = None
+_nvjitlink = None
+_nvjitlink_version = None
 _inited = False
 _use_nvjitlink_backend = None
 _nvjitlink_input_types = None
@@ -50,7 +52,9 @@ class Linker:
         Parameters
         ----------
         target_type : ObjectCodeFormatType | str
-            The type of the target output. Must be either "cubin" or "ptx".
+            The type of the target output. Must be "cubin", "ptx", or
+            "ltoir". Linked LTOIR output requires
+            ``link_time_optimization=True`` and nvJitLink 13.3 or newer.
 
         Returns
         -------
@@ -61,6 +65,16 @@ class Linker:
 
             Ensure that input object codes were compiled with appropriate
             flags for linking (e.g., relocatable device code enabled).
+
+            A CUBIN produced with ``incremental=True`` can be passed directly
+            to another :class:`Linker`, but it can still contain unresolved
+            device references and should be finalized before execution.
+
+            ``"ltoir"`` output contains only the LTOIR carried by the inputs.
+            Direct PTX and CUBIN inputs are rejected because they carry no
+            LTOIR. FATBIN, host object, and library inputs are accepted but
+            may carry no LTOIR, in which case they contribute nothing to the
+            output.
         """
     def get_error_log(self) -> str:
         """Get the error log generated by the linker.
@@ -137,6 +151,11 @@ class LinkerOptions:
     link_time_optimization : bool, optional
         Perform link time optimization.
         Default: False.
+    incremental : bool, optional
+        Perform an incremental link. The result can be passed
+        directly to a later :class:`Linker`. Requires nvJitLink 13.2 or newer
+        and is not supported by the driver linker backend.
+        Default: False.
     ptx : bool, optional
         Emit PTX after linking instead of CUBIN; only supported with ``link_time_optimization=True``.
         Default: False.
@@ -216,6 +235,7 @@ class LinkerOptions:
     split_compile_extended: int | None = None
     no_cache: bool | None = None
     numba_debug: bool | None = None
+    incremental: bool | None = None
 
     def __post_init__(self) -> None: ...
     def _prepare_nvjitlink_options(self, as_bytes: bool=False) -> list[bytes] | list[str]: ...
@@ -241,6 +261,10 @@ class LinkerOptions:
             If nvJitLink backend is not available.
         """
 
+def _require_nvjitlink_version(minimum_version: tuple[int, int], feature: str) -> None:
+    """Check that the cached nvJitLink runtime meets a feature's requirement."""
+def _linked_ltoir_output_module():
+    """Return bindings that can retrieve linked LTOIR without a Cython dependency."""
 def _nvjitlink_has_version_symbol(nvjitlink) -> bool: ...
 def _decide_nvjitlink_or_driver() -> bool:
     """Return True if falling back to the cuLink* driver APIs."""

cuda_core/cuda/core/_linker.pyx

@@ -18,6 +18,7 @@ from cuda.bindings cimport cynvjitlink
 
 from ._rt cimport (
     as_cu,
+    as_intptr,
     as_py,
     create_culink_handle,
     create_nvjitlink_handle,
@@ -101,7 +102,9 @@ cdef class Linker:
         Parameters
         ----------
         target_type : ObjectCodeFormatType | str
-            The type of the target output. Must be either "cubin" or "ptx".
+            The type of the target output. Must be "cubin", "ptx", or
+            "ltoir". Linked LTOIR output requires
+            ``link_time_optimization=True`` and nvJitLink 13.3 or newer.
 
         Returns
         -------
@@ -112,6 +115,16 @@ cdef class Linker:
 
             Ensure that input object codes were compiled with appropriate
             flags for linking (e.g., relocatable device code enabled).
+
+            A CUBIN produced with ``incremental=True`` can be passed directly
+            to another :class:`Linker`, but it can still contain unresolved
+            device references and should be finalized before execution.
+
+            ``"ltoir"`` output contains only the LTOIR carried by the inputs.
+            Direct PTX and CUBIN inputs are rejected because they carry no
+            LTOIR. FATBIN, host object, and library inputs are accepted but
+            may carry no LTOIR, in which case they contribute nothing to the
+            output.
         """
         Linker_check_open(self)
         return Linker_link(self, str(target_type))
@@ -256,6 +269,11 @@ class LinkerOptions:
     link_time_optimization : bool, optional
         Perform link time optimization.
         Default: False.
+    incremental : bool, optional
+        Perform an incremental link. The result can be passed
+        directly to a later :class:`Linker`. Requires nvJitLink 13.2 or newer
+        and is not supported by the driver linker backend.
+        Default: False.
     ptx : bool, optional
         Emit PTX after linking instead of CUBIN; only supported with ``link_time_optimization=True``.
         Default: False.
@@ -336,6 +354,7 @@ class LinkerOptions:
     split_compile_extended: int | None = None
     no_cache: bool | None = None
     numba_debug: bool | None = None
+    incremental: bool | None = None
 
     def __post_init__(self) -> None:
         _lazy_init()
@@ -371,6 +390,8 @@ class LinkerOptions:
             options.append("-verbose")
         if self.link_time_optimization:
             options.append("-lto")
+        if self.incremental:
+            options.append("-r")
         if self.ptx:
             options.append("-ptx")
         if self.optimization_level is not None:
@@ -450,6 +471,8 @@ class LinkerOptions:
         if self.link_time_optimization:
             formatted_options.append(1)
             option_keys.append(_driver.CUjit_option.CU_JIT_LTO)
+        if self.incremental:
+            raise ValueError("incremental option is not supported by the driver API")
         if self.ptx:
             raise ValueError("ptx option is not supported by the driver API")
         if self.optimization_level is not None:
@@ -532,8 +555,13 @@ cdef inline int Linker_init(Linker self, tuple object_codes, object options) exc
     cdef void** c_drv_jit_values_ptr
 
     self._options = options = check_or_create_options(LinkerOptions, options, "Linker options")
+    self._has_ptx_or_cubin_input = False
+    if options.incremental and options.ptx:
+        raise ValueError("incremental and ptx output options cannot be used together")
 
     if _use_nvjitlink_backend:
+        if options.incremental:
+            _require_nvjitlink_version((13, 2), "incremental linking")
         self._use_nvjitlink = True
         options_bytes = options._prepare_nvjitlink_options(as_bytes=True)
         c_num_opts = len(options_bytes)
@@ -639,17 +667,34 @@ cdef inline void Linker_add_code_object(Linker self, object object_code) except
             Linker_annotate_error_log(self, e)
             raise
 
+    if object_code.code_type in ("ptx", "cubin"):
+        self._has_ptx_or_cubin_input = True
+
 
 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"):
+    if target_type not in ("cubin", "ptx", "ltoir"):
         raise ValueError(f"Unsupported target type: {target_type}")
+    if self._options.incremental and target_type == "ptx":
+        raise ValueError("PTX output is not supported for incremental linking")
+    if target_type == "ltoir":
+        if not self._use_nvjitlink:
+            raise ValueError("LTOIR output is not supported by the driver API")
+        if not self._options.link_time_optimization:
+            raise ValueError("LTOIR output requires link_time_optimization=True")
+        if self._has_ptx_or_cubin_input:
+            raise ValueError(
+                'LTOIR output is not supported with "ptx" or "cubin" inputs; '
+                "they carry no LTOIR and would be omitted from the output"
+            )
+        nvjitlink_module = _linked_ltoir_output_module()
 
     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
+    cdef intptr_t c_handle
 
     if self._use_nvjitlink:
         c_nvjitlink_h = as_cu(self._nvjitlink_handle)
@@ -663,14 +708,19 @@ cdef inline object Linker_link(Linker self, str target_type):
             with nogil:
                 HANDLE_RETURN_NVJITLINK(c_nvjitlink_h,
                     cynvjitlink.nvJitLinkGetLinkedCubin(c_nvjitlink_h, c_code_ptr))
-        else:
+        elif target_type == "ptx":
             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_handle = as_intptr(self._nvjitlink_handle)
+            output_size = nvjitlink_module.get_linked_ltoir_size(c_handle)
+            code = bytearray(output_size)
+            nvjitlink_module.get_linked_ltoir(c_handle, code)
     else:
         c_culink_state = as_cu(self._culink_handle)
         try:
@@ -702,6 +752,8 @@ cdef inline void Linker_annotate_error_log(Linker self, object e):
 
 # TODO: revisit this treatment for py313t builds
 _driver = None  # populated if nvJitLink cannot be used
+_nvjitlink = None  # populated if nvJitLink can be used
+_nvjitlink_version = None
 _inited = False
 _use_nvjitlink_backend = None  # set by _decide_nvjitlink_or_driver()
 
@@ -710,6 +762,31 @@ _nvjitlink_input_types = None
 _driver_input_types = None
 
 
+def _require_nvjitlink_version(minimum_version: tuple[int, int], feature: str) -> None:
+    """Check that the cached nvJitLink runtime meets a feature's requirement."""
+    if _nvjitlink_version < minimum_version:
+        required = ".".join(str(component) for component in minimum_version)
+        detected = ".".join(str(component) for component in _nvjitlink_version)
+        raise RuntimeError(f"{feature} requires nvJitLink {required} or newer; found {detected}")
+
+
+# TODO(#2783): Replace this Python-level dispatch with direct cimports once
+# the cuda-bindings runtime floor includes the linked-LTOIR getters.
+def _linked_ltoir_output_module():
+    """Return bindings that can retrieve linked LTOIR without a Cython dependency."""
+    _require_nvjitlink_version((13, 3), "LTOIR output")
+    missing = [
+        name
+        for name in ("get_linked_ltoir_size", "get_linked_ltoir")
+        if not hasattr(_nvjitlink, name)
+    ]
+    if missing:
+        raise RuntimeError(
+            "LTOIR output requires cuda-bindings with " + " and ".join(missing)
+        )
+    return _nvjitlink
+
+
 def _nvjitlink_has_version_symbol(nvjitlink) -> bool:
     # This condition is equivalent to testing for version >= 12.3
     return bool(nvjitlink._inspect_function_pointer("__nvJitLinkVersion"))
@@ -718,10 +795,13 @@ def _nvjitlink_has_version_symbol(nvjitlink) -> bool:
 # 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
+    global _driver, _nvjitlink, _nvjitlink_version, _use_nvjitlink_backend
     if _use_nvjitlink_backend is not None:
         return not _use_nvjitlink_backend
 
+    _nvjitlink = None
+    _nvjitlink_version = None
+
     warn_txt_common = (
         "the driver APIs will be used instead, which do not support"
         " minor version compatibility or linking LTO IRs."
@@ -742,6 +822,9 @@ def _decide_nvjitlink_or_driver() -> bool:
             )
         else:
             if has_version_symbol:
+                detected_version = nvjitlink_module.version()
+                _nvjitlink = nvjitlink_module
+                _nvjitlink_version = detected_version
                 _use_nvjitlink_backend = True
                 return False  # Use nvjitlink
             warn_txt = (

cuda_core/cuda/core/_module.pyi

@@ -355,13 +355,13 @@ class ObjectCode:
         """
     @staticmethod
     def from_object(module: bytes | str, *, name: str='', symbol_mapping: dict[str, str] | None=None) -> ObjectCode:
-        """Create an :class:`ObjectCode` instance from an existing object code.
+        """Create an :class:`ObjectCode` instance from a host object containing device code.
 
         Parameters
         ----------
         module : bytes | str
-            Either a bytes object containing the in-memory object code to load, or
-            a file path string pointing to the on-disk object code to load.
+            Either a bytes object containing the in-memory host object to load, or
+            a file path string pointing to the on-disk host object to load.
         name : str | None
             A human-readable identifier representing this code object.
         symbol_mapping : dict | None

cuda_core/cuda/core/_module.pyx

@@ -741,13 +741,13 @@ cdef class ObjectCode:
 
     @staticmethod
     def from_object(module: bytes | str, *, name: str = "", symbol_mapping: dict[str, str] | None = None) -> ObjectCode:
-        """Create an :class:`ObjectCode` instance from an existing object code.
+        """Create an :class:`ObjectCode` instance from a host object containing device code.
 
         Parameters
         ----------
         module : bytes | str
-            Either a bytes object containing the in-memory object code to load, or
-            a file path string pointing to the on-disk object code to load.
+            Either a bytes object containing the in-memory host object to load, or
+            a file path string pointing to the on-disk host object to load.
         name : str | None
             A human-readable identifier representing this code object.
         symbol_mapping : dict | None

cuda_core/cuda/core/typing.py

@@ -84,7 +84,7 @@ class ObjectCodeFormatType(StrEnum):
     * ``CUBIN`` — device-native CUDA binary.
     * ``LTOIR`` — LTO (link-time optimization) IR for later linking.
     * ``FATBIN`` — fat binary bundling multiple device images.
-    * ``OBJECT`` — relocatable device object.
+    * ``OBJECT`` — host object containing device code.
     * ``LIBRARY`` — device code library.
     """
 

cuda_core/docs/source/release/1.3.0-notes.rst

@@ -9,6 +9,15 @@
 New features
 ------------
 
+- Added incremental linking to :class:`Linker` through
+  ``LinkerOptions(incremental=True)`` on nvJitLink 13.2 or newer. A partial
+  native result is returned as a CUBIN and can be passed directly to another
+  linker. :meth:`Linker.link` also accepts ``"ltoir"`` on nvJitLink 13.3 or
+  newer so incremental LTO chains can retain IR until their final link. Direct
+  PTX and CUBIN inputs are rejected for this target because they carry no LTOIR
+  and nvJitLink would otherwise silently omit them.
+  (`#2369 <https://github.com/NVIDIA/cuda-python/issues/2369>`__)
+
 - Added the read-only :attr:`ManagedBuffer.last_prefetch_location` property,
   which reports the destination requested by the most recent explicit
   prefetch across the buffer. It returns a :class:`Device` or :class:`Host`,