Merge path_finder_dev into main

cuda_bindings/cuda/bindings/_path_finder/README.md

@@ -0,0 +1,55 @@
+# `cuda.bindings.path_finder` Module
+
+## Public API (Work in Progress)
+
+Currently exposes two primary interfaces:
+
+```
+cuda.bindings.path_finder._SUPPORTED_LIBNAMES  # ('nvJitLink', 'nvrtc', 'nvvm')
+cuda.bindings.path_finder._load_nvidia_dynamic_library(libname: str) -> LoadedDL
+```
+
+**Note:**
+These APIs are prefixed with an underscore because they are considered
+experimental while undergoing active development, although already
+reasonably well-tested through CI pipelines.
+
+## Library Loading Search Priority
+
+The `load_nvidia_dynamic_library()` function implements a hierarchical search
+strategy for locating NVIDIA shared libraries:
+
+0. **Check if a library was loaded into the process already by some other means.**
+   - If yes, there is no alternative to skipping the rest of the search logic.
+     The absolute path of the already loaded library will be returned, along
+     with the handle to the library.
+
+1. **NVIDIA Python wheels**
+   - Scans all site-packages to find libraries installed via NVIDIA Python wheels.
+
+2. **OS default mechanisms / Conda environments**
+   - Falls back to native loader:
+     - `dlopen()` on Linux
+     - `LoadLibraryW()` on Windows
+   - CTK installations with system config updates are expected to be discovered:
+     - Linux: Via `/etc/ld.so.conf.d/*cuda*.conf`
+     - Windows: Via `C:\Program Files\NVIDIA GPU Computing Toolkit\CUDA\vX.Y\bin` on system `PATH`
+   - Conda installations are expected to be discovered:
+     - Linux: Via `$ORIGIN/../lib` on `RPATH` (of the `python` binary)
+     - Windows: Via `%CONDA_PREFIX%\Library\bin` on system `PATH`
+
+3. **Environment variables**
+   - Relies on `CUDA_HOME` or `CUDA_PATH` environment variables if set
+     (in that order).
+
+Note that the search is done on a per-library basis. There is no centralized
+mechanism that ensures all libraries are found in the same way.
+
+## Maintenance Requirements
+
+These key components must be updated for new CUDA Toolkit releases:
+
+- `supported_libs.SUPPORTED_LIBNAMES`
+- `supported_libs.SUPPORTED_WINDOWS_DLLS`
+- `supported_libs.SUPPORTED_LINUX_SONAMES`
+- `supported_libs.EXPECTED_LIB_SYMBOLS`

cuda_bindings/cuda/bindings/_path_finder/find_nvidia_dynamic_library.py

@@ -0,0 +1,153 @@
+# Copyright 2024-2025 NVIDIA Corporation.  All rights reserved.
+# SPDX-License-Identifier: LicenseRef-NVIDIA-SOFTWARE-LICENSE
+
+import functools
+import glob
+import os
+
+from cuda.bindings._path_finder.find_sub_dirs import find_sub_dirs_all_sitepackages
+from cuda.bindings._path_finder.supported_libs import IS_WINDOWS, is_suppressed_dll_file
+
+
+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 find_sub_dirs_all_sitepackages(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):
+    nvidia_sub_dirs = ("nvidia", "*", "nvvm", "lib64") if libname == "nvvm" else ("nvidia", "*", "lib")
+    file_wild = so_basename + "*"
+    for lib_dir in find_sub_dirs_all_sitepackages(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, lib_searched_for, error_messages, attachments):
+    nvidia_sub_dirs = ("nvidia", "*", "nvvm", "bin") if libname == "nvvm" else ("nvidia", "*", "bin")
+    for bin_dir in find_sub_dirs_all_sitepackages(nvidia_sub_dirs):
+        dll_name = _find_dll_under_dir(bin_dir, lib_searched_for)
+        if dll_name is not None:
+            return dll_name
+    _no_such_file_in_sub_dirs(nvidia_sub_dirs, lib_searched_for, error_messages, attachments)
+    return None
+
+
+def _get_cuda_home():
+    cuda_home = os.environ.get("CUDA_HOME")
+    if cuda_home is None:
+        cuda_home = os.environ.get("CUDA_PATH")
+    return cuda_home
+
+
+def _find_lib_dir_using_cuda_home(libname):
+    cuda_home = _get_cuda_home()
+    if cuda_home is None:
+        return None
+    if IS_WINDOWS:
+        subdirs = (os.path.join("nvvm", "bin"),) if libname == "nvvm" else ("bin",)
+    else:
+        subdirs = (
+            (os.path.join("nvvm", "lib64"),)
+            if libname == "nvvm"
+            else (
+                "lib64",  # CTK
+                "lib",  # Conda
+            )
+        )
+    for subdir in subdirs:
+        dirname = os.path.join(cuda_home, subdir)
+        if os.path.isdir(dirname):
+            return dirname
+    return None
+
+
+def _find_so_using_lib_dir(lib_dir, so_basename, error_messages, attachments):
+    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}")
+    attachments.append(f'  listdir("{lib_dir}"):')
+    if not os.path.isdir(lib_dir):
+        attachments.append("    DIRECTORY DOES NOT EXIST")
+    else:
+        for node in sorted(os.listdir(lib_dir)):
+            attachments.append(f"    {node}")
+    return None
+
+
+def _find_dll_using_lib_dir(lib_dir, libname, error_messages, attachments):
+    file_wild = libname + "*.dll"
+    dll_name = _find_dll_under_dir(lib_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("{lib_dir}"):')
+    for node in sorted(os.listdir(lib_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 IS_WINDOWS:
+            self.lib_searched_for = f"{libname}*.dll"
+            if self.abs_path is None:
+                self.abs_path = _find_dll_using_nvidia_bin_dirs(
+                    libname, self.lib_searched_for, self.error_messages, self.attachments
+                )
+        else:
+            self.lib_searched_for = f"lib{libname}.so"
+            if self.abs_path is None:
+                self.abs_path = _find_so_using_nvidia_lib_dirs(
+                    libname, self.lib_searched_for, self.error_messages, self.attachments
+                )
+
+    def retry_with_cuda_home_priority_last(self):
+        cuda_home_lib_dir = _find_lib_dir_using_cuda_home(self.libname)
+        if cuda_home_lib_dir is not None:
+            if IS_WINDOWS:
+                self.abs_path = _find_dll_using_lib_dir(
+                    cuda_home_lib_dir, self.libname, self.error_messages, self.attachments
+                )
+            else:
+                self.abs_path = _find_so_using_lib_dir(
+                    cuda_home_lib_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/find_sub_dirs.py

@@ -0,0 +1,52 @@
+# Copyright 2024-2025 NVIDIA Corporation.  All rights reserved.
+# SPDX-License-Identifier: LicenseRef-NVIDIA-SOFTWARE-LICENSE
+
+import functools
+import os
+import site
+import sys
+
+
+def find_sub_dirs_no_cache(parent_dirs, sub_dirs):
+    results = []
+    for base in parent_dirs:
+        stack = [(base, 0)]  # (current_path, index into sub_dirs)
+        while stack:
+            current_path, idx = stack.pop()
+            if idx == len(sub_dirs):
+                if os.path.isdir(current_path):
+                    results.append(current_path)
+                continue
+
+            sub = sub_dirs[idx]
+            if sub == "*":
+                try:
+                    entries = sorted(os.listdir(current_path))
+                except OSError:
+                    continue
+                for entry in entries:
+                    entry_path = os.path.join(current_path, entry)
+                    if os.path.isdir(entry_path):
+                        stack.append((entry_path, idx + 1))
+            else:
+                next_path = os.path.join(current_path, sub)
+                if os.path.isdir(next_path):
+                    stack.append((next_path, idx + 1))
+    return results
+
+
+@functools.cache
+def find_sub_dirs_cached(parent_dirs, sub_dirs):
+    return find_sub_dirs_no_cache(parent_dirs, sub_dirs)
+
+
+def find_sub_dirs(parent_dirs, sub_dirs):
+    return find_sub_dirs_cached(tuple(parent_dirs), tuple(sub_dirs))
+
+
+def find_sub_dirs_sys_path(sub_dirs):
+    return find_sub_dirs(sys.path, sub_dirs)
+
+
+def find_sub_dirs_all_sitepackages(sub_dirs):
+    return find_sub_dirs((site.getusersitepackages(),) + tuple(site.getsitepackages()), sub_dirs)

cuda_bindings/cuda/bindings/_path_finder/load_dl_common.py

@@ -0,0 +1,36 @@
+# Copyright 2025 NVIDIA Corporation.  All rights reserved.
+# SPDX-License-Identifier: LicenseRef-NVIDIA-SOFTWARE-LICENSE
+
+from dataclasses import dataclass
+from typing import Callable, Optional
+
+from cuda.bindings._path_finder.supported_libs import DIRECT_DEPENDENCIES, IS_WINDOWS
+
+if IS_WINDOWS:
+    import pywintypes
+
+    HandleType = pywintypes.HANDLE
+else:
+    HandleType = int
+
+
+@dataclass
+class LoadedDL:
+    handle: HandleType
+    abs_path: Optional[str]
+    was_already_loaded_from_elsewhere: bool
+
+
+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)