Add find_nvidia_dynamic_library() and use from nvjitlink_linux.pyx, nvvm_linux.pyx

Author: rwgk

cuda_bindings/cuda/bindings/_internal/nvjitlink_linux.pyx

@@ -8,8 +8,6 @@ import os
 
 from libc.stdint cimport intptr_t
 
-from .utils cimport get_nvjitlink_dso_version_suffix
-
 from .utils import FunctionNotFoundError, NotSupportedError
 
 from cuda.bindings import path_finder
@@ -57,49 +55,12 @@ cdef void* __nvJitLinkVersion = NULL
 
 
 cdef void* load_library(const int driver_ver) except* with gil:
-    # TODO(rwgk): Move the version check here.
-    # Intentionally ignoring returned value:
-    get_nvjitlink_dso_version_suffix(driver_ver)
-
-    so_basename = "libnvJitLink.so"
-    cdef void* handle = NULL;
-    paths = path_finder.get_cuda_paths()
-    paths_cudalib_dir = paths["cudalib_dir"]
-    if not paths_cudalib_dir:
-        raise RuntimeError("Failure obtaining paths_cudalib_dir")
-    if not paths_cudalib_dir.info:
-        raise RuntimeError("Failure obtaining paths_cudalib_dir.info")
-    primary_so_dir = paths_cudalib_dir.info + "/"
-    candidate_so_dirs = [primary_so_dir]
-    libs = ["/lib/", "/lib64/"]
-    for _ in range(2):
-        alt_dir = libs[0].join(primary_so_dir.rsplit(libs[1], 1))
-        if alt_dir not in candidate_so_dirs:
-            candidate_so_dirs.append(alt_dir)
-        libs.reverse()
-    candidate_so_names = [
-        so_dirname + so_basename
-        for so_dirname in candidate_so_dirs]
-    error_messages = []
-    for so_name in candidate_so_names:
-        if not os.path.exists(so_name):
-            error_messages.append(f"No such file: {so_name}")
-        else:
-            handle = dlopen(so_name.encode(), RTLD_NOW | RTLD_GLOBAL)
-            if handle != NULL:
-                return handle
-            err_msg = dlerror().decode(errors="backslashreplace")
-            error_messages.append(f"Failed to dlopen {so_name}: {err_msg}")
-    attachment = []
-    for so_dirname in candidate_so_dirs:
-        attachment.append(f"  listdir({repr(so_dirname)}):")
-        if not os.path.isdir(so_dirname):
-            attachment.append("    DIRECTORY DOES NOT EXIST")
-        else:
-            for node in sorted(os.listdir(so_dirname)):
-                attachment.append(f"    {node}")
-    attachment = "\n".join(attachment)
-    raise RuntimeError(f"Unable to load {so_basename} from: {', '.join(error_messages)}\n{attachment}")
+    so_name = path_finder.find_nvidia_dynamic_library("nvJitLink")
+    cdef void* handle = dlopen(so_name.encode(), RTLD_NOW | RTLD_GLOBAL)
+    if handle != NULL:
+        return handle
+    err_msg = dlerror().decode(errors="backslashreplace")
+    raise RuntimeError(f"Failed to dlopen {so_name}: {err_msg}")
 
 
 cdef int _check_or_init_nvjitlink() except -1 nogil:

cuda_bindings/cuda/bindings/_internal/nvvm_linux.pyx

@@ -51,17 +51,12 @@ cdef void* __nvvmGetProgramLog = NULL
 
 
 cdef void* load_library(const int driver_ver) except* with gil:
-    cdef void* handle = NULL;
-    paths = path_finder.get_cuda_paths()
-    paths_nvvm = paths["nvvm"]
-    if paths_nvvm:
-        so_name = paths_nvvm.info
-        handle = dlopen(so_name.encode(), RTLD_NOW | RTLD_GLOBAL)
-        if handle == NULL:
-            err_msg = dlerror()
-            raise RuntimeError(f'Failed to dlopen {so_name} ({err_msg.decode()})')
+    so_name = path_finder.find_nvidia_dynamic_library("nvvm")
+    cdef void* handle = dlopen(so_name.encode(), RTLD_NOW | RTLD_GLOBAL)
+    if handle != NULL:
         return handle
-    raise RuntimeError('Unable to locate libnvvm.so')
+    err_msg = dlerror().decode(errors="backslashreplace")
+    raise RuntimeError(f"Failed to dlopen {so_name}: {err_msg}")
 
 
 cdef int _check_or_init_nvvm() except -1 nogil:

cuda_bindings/cuda/bindings/_path_finder_utils/find_nvidia_dynamic_library.py

@@ -0,0 +1,76 @@
+# Copyright 2024-2025 NVIDIA Corporation.  All rights reserved.
+#
+# SPDX-License-Identifier: LicenseRef-NVIDIA-SOFTWARE-LICENSE
+
+import functools
+import os
+
+from .cuda_paths import get_cuda_paths
+from .find_nvidia_lib_dirs import find_nvidia_lib_dirs
+
+
+def _find_using_nvidia_lib_dirs(so_basename, error_messages, attachments):
+    for lib_dir in find_nvidia_lib_dirs():
+        so_name = os.path.join(lib_dir, so_basename)
+        if os.path.isfile(so_name):
+            return so_name
+        error_messages.append(f"No such file: {so_name}")
+    for lib_dir in find_nvidia_lib_dirs():
+        attachments.append(f"  listdir({repr(lib_dir)}):")
+        for node in sorted(os.listdir(lib_dir)):
+            attachments.append(f"    {node}")
+    return None
+
+
+def _get_cuda_paths_info(key, error_messages):
+    env_path_tuple = get_cuda_paths()[key]
+    if not env_path_tuple:
+        error_messages.append(f'Failure obtaining get_cuda_paths()["{key}"]')
+        return None
+    if not env_path_tuple.info:
+        error_messages.append(f'Failure obtaining get_cuda_paths()["{key}"].info')
+        return None
+    return env_path_tuple.info
+
+
+def _find_using_lib_dir(so_basename, error_messages, attachments):
+    lib_dir = _get_cuda_paths_info("cudalib_dir", error_messages)
+    primary_so_dir = lib_dir + "/"
+    candidate_so_dirs = [primary_so_dir]
+    libs = ["/lib/", "/lib64/"]
+    for _ in range(2):
+        alt_dir = libs[0].join(primary_so_dir.rsplit(libs[1], 1))
+        if alt_dir not in candidate_so_dirs:
+            candidate_so_dirs.append(alt_dir)
+        libs.reverse()
+    candidate_so_names = [so_dirname + so_basename for so_dirname in candidate_so_dirs]
+    error_messages = []
+    for so_name in candidate_so_names:
+        if os.path.isfile(so_name):
+            return so_name
+        error_messages.append(f"No such file: {so_name}")
+    for so_dirname in candidate_so_dirs:
+        attachments.append(f"  listdir({repr(so_dirname)}):")
+        if not os.path.isdir(so_dirname):
+            attachments.append("    DIRECTORY DOES NOT EXIST")
+        else:
+            for node in sorted(os.listdir(so_dirname)):
+                attachments.append(f"    {node}")
+    return None
+
+
+@functools.cache
+def find_nvidia_dynamic_library(libbasename):
+    so_basename = f"lib{libbasename}.so"
+    error_messages = []
+    attachments = []
+    so_name = _find_using_nvidia_lib_dirs(so_basename, error_messages, attachments)
+    if so_name is None:
+        if libbasename == "nvvm":
+            so_name = _get_cuda_paths_info("nvvm", error_messages)
+        else:
+            so_name = _find_using_lib_dir(so_basename, error_messages, attachments)
+    if so_name is None:
+        attachments = "\n".join(attachments)
+        raise RuntimeError(f"Unable to load {so_basename} from: {', '.join(error_messages)}\n{attachments}")
+    return so_name

cuda_bindings/cuda/bindings/_path_finder_utils/find_nvidia_lib_dirs.py

@@ -0,0 +1,30 @@
+# Copyright 2024-2025 NVIDIA Corporation.  All rights reserved.
+#
+# SPDX-License-Identifier: LicenseRef-NVIDIA-SOFTWARE-LICENSE
+
+import functools
+import os
+import sys
+
+
+@functools.cache
+def _find_nvidia_lib_dirs(sys_path):
+    results = []
+    for base in sys_path:
+        nvidia_base = os.path.join(base, "nvidia")
+        if not os.path.isdir(nvidia_base):
+            continue
+        try:
+            subdirs = os.listdir(nvidia_base)
+        except OSError:
+            continue
+        for sub in subdirs:
+            sub_path = os.path.join(nvidia_base, sub)
+            lib_path = os.path.join(sub_path, "lib")
+            if os.path.isdir(lib_path):
+                results.append(lib_path)
+    return results
+
+
+def find_nvidia_lib_dirs():
+    return _find_nvidia_lib_dirs(tuple(sys.path))

cuda_bindings/cuda/bindings/path_finder.py

@@ -16,8 +16,10 @@
     get_nvidia_static_cudalib_ctk,
     get_system_ctk,
 )
+from cuda.bindings._path_finder_utils.find_nvidia_dynamic_library import find_nvidia_dynamic_library
 
 __all__ = [
+    "find_nvidia_dynamic_library",
     "get_conda_ctk",
     "get_conda_include_dir",
     "get_cuda_home",

cuda_bindings/tests/path_finder.py

@@ -4,3 +4,6 @@
 
 for k, v in paths.items():
     print(f"{k}: {v}", flush=True)
+
+print(path_finder.find_nvidia_dynamic_library("nvvm"))
+print(path_finder.find_nvidia_dynamic_library("nvJitLink"))