Skip to content

Commit 7c65677

Browse files
rwgkleofang
andauthored
_decide_nvjitlink_or_driver() bug fix and enhancements (NVIDIA#1055)
* _decide_nvjitlink_or_driver(): catch RuntimeError (bug fix), use importlib + ModuleNotFoundError (more selective than ImportError) and produce specific error messages * Fix misunderstanding: RuntimeError is raised only from inner_nvjitlink._inspect_function_pointer() * Better way of formatting warning messages. * Change from importlib.import_module() to plain import (the latter does also raise ModuleNotFoundError) * Enhance to warning messages, to make them actionable. * Factor out _nvjitlink_has_version_symbol() for clarity and testability This aids unit testing by allowing localized stubbing of the version-symbol check, without needing to patch the full inner nvjitlink module. * Add test_linker_warnings.py As generated by ChatGPT 5, with minor manual tweaks. * Fix "the the" oversight * Replace "culink APIs" → "driver APIs" in warning message. * Fix oversight: test_linker_warnings.py needs to be updated after commit 0948942 * fix skipping the check for nvidia-smi (NVIDIA#1084) * rm cuda_core/tests/test_linker_warnings.py: see NVIDIA#1095 --------- Co-authored-by: Leo Fang <leof@nvidia.com>
1 parent 4cbb627 commit 7c65677

1 file changed

Lines changed: 33 additions & 19 deletions

File tree

cuda_core/cuda/core/experimental/_linker.py

Lines changed: 33 additions & 19 deletions
Original file line numberDiff line numberDiff line change
@@ -5,6 +5,7 @@
55
from __future__ import annotations
66

77
import ctypes
8+
import sys
89
import weakref
910
from contextlib import contextmanager
1011
from dataclasses import dataclass
@@ -28,6 +29,11 @@
2829
_nvjitlink_input_types = None # populated if nvJitLink cannot be used
2930

3031

32+
def _nvjitlink_has_version_symbol(inner_nvjitlink) -> bool:
33+
# This condition is equivalent to testing for version >= 12.3
34+
return bool(inner_nvjitlink._inspect_function_pointer("__nvJitLinkVersion"))
35+
36+
3137
# Note: this function is reused in the tests
3238
def _decide_nvjitlink_or_driver() -> bool:
3339
"""Returns True if falling back to the cuLink* driver APIs."""
@@ -37,28 +43,36 @@ def _decide_nvjitlink_or_driver() -> bool:
3743

3844
_driver_ver = handle_return(driver.cuDriverGetVersion())
3945
_driver_ver = (_driver_ver // 1000, (_driver_ver % 1000) // 10)
46+
47+
warn_txt_common = (
48+
"the driver APIs will be used instead, which do not support"
49+
" minor version compatibility or linking LTO IRs."
50+
" For best results, consider upgrading to a recent version of"
51+
)
52+
4053
try:
41-
from cuda.bindings import nvjitlink as _nvjitlink
42-
from cuda.bindings._internal import nvjitlink as inner_nvjitlink
43-
except ImportError:
44-
# binding is not available
45-
_nvjitlink = None
54+
import cuda.bindings.nvjitlink as _nvjitlink
55+
except ModuleNotFoundError:
56+
warn_txt = f"cuda.bindings.nvjitlink is not available, therefore {warn_txt_common} cuda-bindings."
4657
else:
47-
if inner_nvjitlink._inspect_function_pointer("__nvJitLinkVersion") == 0:
48-
# binding is available, but nvJitLink is not installed
49-
_nvjitlink = None
50-
51-
if _nvjitlink is None:
52-
warn(
53-
"nvJitLink is not installed or too old (<12.3). Therefore it is not usable "
54-
"and the culink APIs will be used instead.",
55-
stacklevel=3,
56-
category=RuntimeWarning,
58+
from cuda.bindings._internal import nvjitlink as inner_nvjitlink
59+
60+
try:
61+
if _nvjitlink_has_version_symbol(inner_nvjitlink):
62+
return False # Use nvjitlink
63+
except RuntimeError:
64+
warn_detail = "not available"
65+
else:
66+
warn_detail = "too old (<12.3)"
67+
warn_txt = (
68+
f"{'nvJitLink*.dll' if sys.platform == 'win32' else 'libnvJitLink.so*'} is {warn_detail}."
69+
f" Therefore cuda.bindings.nvjitlink is not usable and {warn_txt_common} nvJitLink."
5770
)
58-
_driver = driver
59-
return True
60-
else:
61-
return False
71+
_nvjitlink = None
72+
73+
warn(warn_txt, stacklevel=2, category=RuntimeWarning)
74+
_driver = driver
75+
return True
6276

6377

6478
def _lazy_init():

0 commit comments

Comments
 (0)