First version of cuda.bindings.path_finder

cuda_bindings/cuda/bindings/_path_finder/find_nvidia_dynamic_library.py

@@ -0,0 +1,156 @@
+# Copyright 2024-2025 NVIDIA Corporation.  All rights reserved.
+# SPDX-License-Identifier: LicenseRef-NVIDIA-SOFTWARE-LICENSE
+
+import functools
+import glob
+import os
+import sys
+
+from .cuda_paths import get_cuda_paths
+from .supported_libs import is_suppressed_dll_file
+from .sys_path_find_sub_dirs import sys_path_find_sub_dirs
+
+
+def _no_such_file_in_sub_dirs(sub_dirs, file_wild, error_messages, attachments):
+    error_messages.append(f"No such file: {file_wild}")
+    for sub_dir in sys_path_find_sub_dirs(sub_dirs):
+        attachments.append(f'  listdir("{sub_dir}"):')
+        for node in sorted(os.listdir(sub_dir)):
+            attachments.append(f"    {node}")
+
+
+def _find_so_using_nvidia_lib_dirs(libname, so_basename, error_messages, attachments):
+    if libname == "nvvm":  # noqa: SIM108
+        nvidia_sub_dirs = ("nvidia", "*", "nvvm", "lib64")
+    else:
+        nvidia_sub_dirs = ("nvidia", "*", "lib")
+    file_wild = so_basename + "*"
+    for lib_dir in sys_path_find_sub_dirs(nvidia_sub_dirs):
+        # First look for an exact match
+        so_name = os.path.join(lib_dir, so_basename)
+        if os.path.isfile(so_name):
+            return so_name
+        # Look for a versioned library
+        # Using sort here mainly to make the result deterministic.
+        for so_name in sorted(glob.glob(os.path.join(lib_dir, file_wild))):
+            if os.path.isfile(so_name):
+                return so_name
+    _no_such_file_in_sub_dirs(nvidia_sub_dirs, file_wild, error_messages, attachments)
+    return None
+
+
+def _find_dll_under_dir(dirpath, file_wild):
+    for path in sorted(glob.glob(os.path.join(dirpath, file_wild))):
+        if not os.path.isfile(path):
+            continue
+        if not is_suppressed_dll_file(os.path.basename(path)):
+            return path
+    return None
+
+
+def _find_dll_using_nvidia_bin_dirs(libname, error_messages, attachments):
+    if libname == "nvvm":  # noqa: SIM108
+        nvidia_sub_dirs = ("nvidia", "*", "nvvm", "bin")
+    else:
+        nvidia_sub_dirs = ("nvidia", "*", "bin")
+    file_wild = libname + "*.dll"
+    for bin_dir in sys_path_find_sub_dirs(nvidia_sub_dirs):
+        dll_name = _find_dll_under_dir(bin_dir, file_wild)
+        if dll_name is not None:
+            return dll_name
+    _no_such_file_in_sub_dirs(nvidia_sub_dirs, file_wild, error_messages, attachments)
+    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_so_using_cudalib_dir(so_basename, error_messages, attachments):
+    cudalib_dir = _get_cuda_paths_info("cudalib_dir", error_messages)
+    if cudalib_dir is None:
+        return None
+    primary_so_dir = cudalib_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]
+    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("{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
+
+
+def _find_dll_using_cudalib_dir(libname, error_messages, attachments):
+    cudalib_dir = _get_cuda_paths_info("cudalib_dir", error_messages)
+    if cudalib_dir is None:
+        return None
+    file_wild = libname + "*.dll"
+    dll_name = _find_dll_under_dir(cudalib_dir, file_wild)
+    if dll_name is not None:
+        return dll_name
+    error_messages.append(f"No such file: {file_wild}")
+    attachments.append(f'  listdir("{cudalib_dir}"):')
+    for node in sorted(os.listdir(cudalib_dir)):
+        attachments.append(f"    {node}")
+    return None
+
+
+class _find_nvidia_dynamic_library:
+    def __init__(self, libname: str):
+        self.libname = libname
+        self.error_messages = []
+        self.attachments = []
+        self.abs_path = None
+
+        if sys.platform == "win32":
+            self.abs_path = _find_dll_using_nvidia_bin_dirs(libname, self.error_messages, self.attachments)
+            if self.abs_path is None:
+                if libname == "nvvm":
+                    self.abs_path = _get_cuda_paths_info("nvvm", self.error_messages)
+                else:
+                    self.abs_path = _find_dll_using_cudalib_dir(libname, self.error_messages, self.attachments)
+            self.lib_searched_for = f"{libname}*.dll"
+        else:
+            self.lib_searched_for = f"lib{libname}.so"
+            self.abs_path = _find_so_using_nvidia_lib_dirs(
+                libname, self.lib_searched_for, self.error_messages, self.attachments
+            )
+            if self.abs_path is None:
+                if libname == "nvvm":
+                    self.abs_path = _get_cuda_paths_info("nvvm", self.error_messages)
+                else:
+                    self.abs_path = _find_so_using_cudalib_dir(
+                        self.lib_searched_for, self.error_messages, self.attachments
+                    )
+
+    def raise_if_abs_path_is_None(self):
+        if self.abs_path:
+            return self.abs_path
+        err = ", ".join(self.error_messages)
+        att = "\n".join(self.attachments)
+        raise RuntimeError(f'Failure finding "{self.lib_searched_for}": {err}\n{att}')
+
+
+@functools.cache
+def find_nvidia_dynamic_library(libname: str) -> str:
+    return _find_nvidia_dynamic_library(libname).raise_if_abs_path_is_None()

cuda_bindings/cuda/bindings/_path_finder/findlib.py

@@ -0,0 +1,69 @@
+# Forked from:
+# https://github.com/numba/numba/blob/f0d24824fcd6a454827e3c108882395d00befc04/numba/misc/findlib.py
+
+import os
+import re
+import sys
+
+
+def get_lib_dirs():
+    """
+    Anaconda specific
+    """
+    if sys.platform == "win32":
+        # on windows, historically `DLLs` has been used for CUDA libraries,
+        # since approximately CUDA 9.2, `Library\bin` has been used.
+        dirnames = ["DLLs", os.path.join("Library", "bin")]
+    else:
+        dirnames = [
+            "lib",
+        ]
+    libdirs = [os.path.join(sys.prefix, x) for x in dirnames]
+    return libdirs
+
+
+DLLNAMEMAP = {
+    "linux": r"lib%(name)s\.so\.%(ver)s$",
+    "linux2": r"lib%(name)s\.so\.%(ver)s$",
+    "linux-static": r"lib%(name)s\.a$",
+    "darwin": r"lib%(name)s\.%(ver)s\.dylib$",
+    "win32": r"%(name)s%(ver)s\.dll$",
+    "win32-static": r"%(name)s\.lib$",
+    "bsd": r"lib%(name)s\.so\.%(ver)s$",
+}
+
+RE_VER = r"[0-9]*([_\.][0-9]+)*"
+
+
+def find_lib(libname, libdir=None, platform=None, static=False):
+    platform = platform or sys.platform
+    platform = "bsd" if "bsd" in platform else platform
+    if static:
+        platform = f"{platform}-static"
+    if platform not in DLLNAMEMAP:
+        # Return empty list if platform name is undefined.
+        # Not all platforms define their static library paths.
+        return []
+    pat = DLLNAMEMAP[platform] % {"name": libname, "ver": RE_VER}
+    regex = re.compile(pat)
+    return find_file(regex, libdir)
+
+
+def find_file(pat, libdir=None):
+    if libdir is None:
+        libdirs = get_lib_dirs()
+    elif isinstance(libdir, str):
+        libdirs = [
+            libdir,
+        ]
+    else:
+        libdirs = list(libdir)
+    files = []
+    for ldir in libdirs:
+        try:
+            entries = os.listdir(ldir)
+        except FileNotFoundError:
+            continue
+        candidates = [os.path.join(ldir, ent) for ent in entries if pat.match(ent)]
+        files.extend([c for c in candidates if os.path.isfile(c)])
+    return files

cuda_bindings/cuda/bindings/_path_finder/load_dl_common.py

@@ -0,0 +1,60 @@
+# Copyright 2025 NVIDIA Corporation.  All rights reserved.
+# SPDX-License-Identifier: LicenseRef-NVIDIA-SOFTWARE-LICENSE
+
+from dataclasses import dataclass
+from typing import Callable, Optional
+
+from .supported_libs import DIRECT_DEPENDENCIES
+
+
+@dataclass
+class LoadedDL:
+    """Represents a loaded dynamic library.
+
+    Attributes:
+        handle: The library handle (can be converted to void* in Cython)
+        abs_path: The absolute path to the library file
+        was_already_loaded_from_elsewhere: Whether the library was already loaded
+    """
+
+    # ATTENTION: To convert `handle` back to `void*` in cython:
+    #     Linux:   `cdef void* ptr = <void*><uintptr_t>`
+    #     Windows: `cdef void* ptr = <void*><intptr_t>`
+    handle: int
+    abs_path: Optional[str]
+    was_already_loaded_from_elsewhere: bool
+
+
+def add_dll_directory(dll_abs_path: str) -> None:
+    """Add a DLL directory to the search path and update PATH environment variable.
+
+    Args:
+        dll_abs_path: Absolute path to the DLL file
+
+    Raises:
+        AssertionError: If the directory containing the DLL does not exist
+    """
+    import os
+
+    dirpath = os.path.dirname(dll_abs_path)
+    assert os.path.isdir(dirpath), dll_abs_path
+    # Add the DLL directory to the search path
+    os.add_dll_directory(dirpath)
+    # Update PATH as a fallback for dependent DLL resolution
+    curr_path = os.environ.get("PATH")
+    os.environ["PATH"] = dirpath if curr_path is None else os.pathsep.join((curr_path, dirpath))
+
+
+def load_dependencies(libname: str, load_func: Callable[[str], LoadedDL]) -> None:
+    """Load all dependencies for a given library.
+
+    Args:
+        libname: The name of the library whose dependencies should be loaded
+        load_func: The function to use for loading libraries (e.g. load_nvidia_dynamic_library)
+
+    Example:
+        >>> load_dependencies("cudart", load_nvidia_dynamic_library)
+        # This will load all dependencies of cudart using the provided loading function
+    """
+    for dep in DIRECT_DEPENDENCIES.get(libname, ()):
+        load_func(dep)