Save result of factoring out load_dl_common.py, load_dl_linux.py, load_dl_windows.py with the help of Cursor.

rwgk · rwgk · commit b5cef1bc75d1 · 2025-04-25T16:45:10.000-07:00
diff --git a/cuda_bindings/cuda/bindings/_path_finder/load_dl_common.py b/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)
diff --git a/cuda_bindings/cuda/bindings/_path_finder/load_dl_linux.py b/cuda_bindings/cuda/bindings/_path_finder/load_dl_linux.py
@@ -0,0 +1,125 @@
+# Copyright 2025 NVIDIA Corporation.  All rights reserved.
+# SPDX-License-Identifier: LicenseRef-NVIDIA-SOFTWARE-LICENSE
+
+import ctypes
+import ctypes.util
+import os
+from typing import Optional
+
+from .load_dl_common import LoadedDL
+
+CDLL_MODE = os.RTLD_NOW | os.RTLD_GLOBAL
+
+LIBDL_PATH = ctypes.util.find_library("dl") or "libdl.so.2"
+LIBDL = ctypes.CDLL(LIBDL_PATH)
+LIBDL.dladdr.argtypes = [ctypes.c_void_p, ctypes.c_void_p]
+LIBDL.dladdr.restype = ctypes.c_int
+
+
+class Dl_info(ctypes.Structure):
+    """Structure used by dladdr to return information about a loaded symbol."""
+
+    _fields_ = [
+        ("dli_fname", ctypes.c_char_p),  # path to .so
+        ("dli_fbase", ctypes.c_void_p),
+        ("dli_sname", ctypes.c_char_p),
+        ("dli_saddr", ctypes.c_void_p),
+    ]
+
+
+def abs_path_for_dynamic_library(libname: str, handle: int) -> Optional[str]:
+    """Get the absolute path of a loaded dynamic library on Linux.
+
+    Args:
+        libname: The name of the library
+        handle: The library handle
+
+    Returns:
+        The absolute path to the library file, or None if no expected symbol is found
+
+    Raises:
+        OSError: If dladdr fails to get information about the symbol
+    """
+    from .supported_libs import EXPECTED_LIB_SYMBOLS
+
+    for symbol_name in EXPECTED_LIB_SYMBOLS[libname]:
+        symbol = getattr(handle, symbol_name, None)
+        if symbol is not None:
+            break
+    else:
+        return None
+
+    addr = ctypes.cast(symbol, ctypes.c_void_p)
+    info = Dl_info()
+    if LIBDL.dladdr(addr, ctypes.byref(info)) == 0:
+        raise OSError(f"dladdr failed for {libname=!r}")
+    return info.dli_fname.decode()
+
+
+def check_if_already_loaded(libname: str) -> Optional[LoadedDL]:
+    """Check if the library is already loaded in the process.
+
+    Args:
+        libname: The name of the library to check
+
+    Returns:
+        A LoadedDL object if the library is already loaded, None otherwise
+
+    Example:
+        >>> loaded = check_if_already_loaded("cudart")
+        >>> if loaded is not None:
+        ...     print(f"Library already loaded from {loaded.abs_path}")
+    """
+    from .supported_libs import SUPPORTED_LINUX_SONAMES
+
+    for soname in SUPPORTED_LINUX_SONAMES.get(libname, ()):
+        try:
+            handle = ctypes.CDLL(soname, mode=os.RTLD_NOLOAD)
+        except OSError:
+            continue
+        else:
+            return LoadedDL(handle._handle, abs_path_for_dynamic_library(libname, handle), True)
+    return None
+
+
+def load_with_system_search(libname: str, soname: str) -> Optional[LoadedDL]:
+    """Try to load a library using system search paths.
+
+    Args:
+        libname: The name of the library to load
+        soname: The soname to search for
+
+    Returns:
+        A LoadedDL object if successful, None if the library cannot be loaded
+
+    Raises:
+        RuntimeError: If the library is loaded but no expected symbol is found
+    """
+    try:
+        handle = ctypes.CDLL(soname, CDLL_MODE)
+        abs_path = abs_path_for_dynamic_library(libname, handle)
+        if abs_path is None:
+            raise RuntimeError(f"No expected symbol for {libname=!r}")
+        return LoadedDL(handle._handle, abs_path, False)
+    except OSError:
+        return None
+
+
+def load_with_abs_path(libname: str, found_path: str) -> LoadedDL:
+    """Load a dynamic library from the given path.
+
+    Args:
+        libname: The name of the library to load
+        found_path: The absolute path to the library file
+
+    Returns:
+        A LoadedDL object representing the loaded library
+
+    Raises:
+        RuntimeError: If the library cannot be loaded
+    """
+    try:
+        handle = ctypes.CDLL(found_path, CDLL_MODE)
+    except OSError as e:
+        raise RuntimeError(f"Failed to dlopen {found_path}: {e}") from e
+    return LoadedDL(handle._handle, found_path, False)
diff --git a/cuda_bindings/cuda/bindings/_path_finder/load_dl_windows.py b/cuda_bindings/cuda/bindings/_path_finder/load_dl_windows.py
@@ -0,0 +1,141 @@
+# Copyright 2025 NVIDIA Corporation.  All rights reserved.
+# SPDX-License-Identifier: LicenseRef-NVIDIA-SOFTWARE-LICENSE
+
+import ctypes
+import ctypes.wintypes
+import functools
+from typing import Optional
+
+import pywintypes
+import win32api
+
+from .load_dl_common import LoadedDL, add_dll_directory
+
+# Mirrors WinBase.h (unfortunately not defined already elsewhere)
+WINBASE_LOAD_LIBRARY_SEARCH_DLL_LOAD_DIR = 0x00000100
+WINBASE_LOAD_LIBRARY_SEARCH_DEFAULT_DIRS = 0x00001000
+
+
+def abs_path_for_dynamic_library(handle: int) -> str:
+    """Get the absolute path of a loaded dynamic library on Windows.
+
+    Args:
+        handle: The library handle
+
+    Returns:
+        The absolute path to the DLL file
+
+    Raises:
+        OSError: If GetModuleFileNameW fails
+    """
+    buf = ctypes.create_unicode_buffer(260)
+    n_chars = ctypes.windll.kernel32.GetModuleFileNameW(ctypes.wintypes.HMODULE(handle), buf, len(buf))
+    if n_chars == 0:
+        raise OSError("GetModuleFileNameW failed")
+    return buf.value
+
+
+@functools.cache
+def cuDriverGetVersion() -> int:
+    """Get the CUDA driver version.
+
+    Returns:
+        The CUDA driver version number
+
+    Raises:
+        AssertionError: If the driver version cannot be obtained
+    """
+    handle = win32api.LoadLibrary("nvcuda.dll")
+
+    kernel32 = ctypes.WinDLL("kernel32", use_last_error=True)
+    GetProcAddress = kernel32.GetProcAddress
+    GetProcAddress.argtypes = [ctypes.wintypes.HMODULE, ctypes.wintypes.LPCSTR]
+    GetProcAddress.restype = ctypes.c_void_p
+    cuDriverGetVersion = GetProcAddress(handle, b"cuDriverGetVersion")
+    assert cuDriverGetVersion
+
+    FUNC_TYPE = ctypes.CFUNCTYPE(ctypes.c_int, ctypes.POINTER(ctypes.c_int))
+    cuDriverGetVersion_fn = FUNC_TYPE(cuDriverGetVersion)
+    driver_ver = ctypes.c_int()
+    err = cuDriverGetVersion_fn(ctypes.byref(driver_ver))
+    assert err == 0
+    return driver_ver.value
+
+
+def check_if_already_loaded(libname: str) -> Optional[LoadedDL]:
+    """Check if the library is already loaded in the process.
+
+    Args:
+        libname: The name of the library to check
+
+    Returns:
+        A LoadedDL object if the library is already loaded, None otherwise
+
+    Example:
+        >>> loaded = check_if_already_loaded("cudart")
+        >>> if loaded is not None:
+        ...     print(f"Library already loaded from {loaded.abs_path}")
+    """
+    from .supported_libs import SUPPORTED_WINDOWS_DLLS
+
+    for dll_name in SUPPORTED_WINDOWS_DLLS.get(libname, ()):
+        try:
+            handle = win32api.GetModuleHandle(dll_name)
+        except pywintypes.error:
+            continue
+        else:
+            return LoadedDL(handle, abs_path_for_dynamic_library(handle), True)
+    return None
+
+
+def load_with_system_search(name: str, _unused: str) -> Optional[LoadedDL]:
+    """Try to load a DLL using system search paths.
+
+    Args:
+        name: The name of the library to load
+        _unused: Unused parameter (kept for interface consistency)
+
+    Returns:
+        A LoadedDL object if successful, None if the library cannot be loaded
+    """
+    from .supported_libs import SUPPORTED_WINDOWS_DLLS
+
+    driver_ver = cuDriverGetVersion()
+    del driver_ver  # Keeping this here because it will probably be needed in the future.
+
+    dll_names = SUPPORTED_WINDOWS_DLLS.get(name)
+    if dll_names is None:
+        return None
+
+    for dll_name in dll_names:
+        handle = ctypes.windll.kernel32.LoadLibraryW(ctypes.c_wchar_p(dll_name))
+        if handle:
+            return LoadedDL(handle, abs_path_for_dynamic_library(handle), False)
+
+    return None
+
+
+def load_with_abs_path(libname: str, found_path: str) -> LoadedDL:
+    """Load a dynamic library from the given path.
+
+    Args:
+        libname: The name of the library to load
+        found_path: The absolute path to the DLL file
+
+    Returns:
+        A LoadedDL object representing the loaded library
+
+    Raises:
+        RuntimeError: If the DLL cannot be loaded
+    """
+    from .supported_libs import LIBNAMES_REQUIRING_OS_ADD_DLL_DIRECTORY
+
+    if libname in LIBNAMES_REQUIRING_OS_ADD_DLL_DIRECTORY:
+        add_dll_directory(found_path)
+
+    flags = WINBASE_LOAD_LIBRARY_SEARCH_DEFAULT_DIRS | WINBASE_LOAD_LIBRARY_SEARCH_DLL_LOAD_DIR
+    try:
+        handle = win32api.LoadLibraryEx(found_path, 0, flags)
+    except pywintypes.error as e:
+        raise RuntimeError(f"Failed to load DLL at {found_path}: {e}") from e
+    return LoadedDL(handle, found_path, False)
diff --git a/cuda_bindings/cuda/bindings/_path_finder/load_nvidia_dynamic_library.py b/cuda_bindings/cuda/bindings/_path_finder/load_nvidia_dynamic_library.py
