Skip to content
Merged
1 change: 1 addition & 0 deletions cuda_core/cuda/core/_linker.pxd
Original file line number Diff line number Diff line change
Expand Up @@ -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()
Expand Down
26 changes: 25 additions & 1 deletion cuda_core/cuda/core/_linker.pyi
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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
-------
Expand All @@ -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.
Expand Down Expand Up @@ -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.
Expand Down Expand Up @@ -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]: ...
Expand All @@ -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."""
Expand Down
91 changes: 87 additions & 4 deletions cuda_core/cuda/core/_linker.pyx
Original file line number Diff line number Diff line change
Expand Up @@ -18,6 +18,7 @@ from cuda.bindings cimport cynvjitlink

from ._resource_handles cimport (
as_cu,
as_intptr,
as_py,
create_culink_handle,
create_nvjitlink_handle,
Expand Down Expand Up @@ -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
-------
Expand All @@ -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))
Expand Down Expand Up @@ -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.
Expand Down Expand Up @@ -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()
Expand Down Expand Up @@ -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:
Expand Down Expand Up @@ -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:
Expand Down Expand Up @@ -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)
Expand Down Expand Up @@ -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":
Comment thread
isVoid marked this conversation as resolved.
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)
Expand All @@ -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:
Expand Down Expand Up @@ -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()

Expand All @@ -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():
Comment thread
isVoid marked this conversation as resolved.
"""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"))
Expand All @@ -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."
Expand All @@ -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 = (
Expand Down
6 changes: 3 additions & 3 deletions cuda_core/cuda/core/_module.pyi
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
6 changes: 3 additions & 3 deletions cuda_core/cuda/core/_module.pyx
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
2 changes: 1 addition & 1 deletion cuda_core/cuda/core/typing.py
Original file line number Diff line number Diff line change
Expand Up @@ -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.
"""

Expand Down
9 changes: 9 additions & 0 deletions cuda_core/docs/source/release/1.3.0-notes.rst
Original file line number Diff line number Diff line change
Expand Up @@ -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`,
Expand Down
Loading
Loading