Coverage for cuda/pathfinder/_dynamic_libs/load_dl_windows.py: 81.43%

70 statements  

« prev     ^ index     » next       coverage.py v7.15.2, created at 2026-07-29 01:38 +0000

1# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. 

2# SPDX-License-Identifier: Apache-2.0 

3 

4from __future__ import annotations 

5 

6import ctypes 

7import ctypes.wintypes 

8import os 

9import struct 

10import warnings 

11from typing import TYPE_CHECKING 

12 

13from cuda.pathfinder._dynamic_libs.load_dl_common import LoadedDL 

14 

15if TYPE_CHECKING: 

16 from cuda.pathfinder._dynamic_libs.lib_descriptor import LibDescriptor 

17 

18# Mirrors WinBase.h (unfortunately not defined already elsewhere) 

19WINBASE_LOAD_LIBRARY_SEARCH_DLL_LOAD_DIR = 0x00000100 

20WINBASE_LOAD_LIBRARY_SEARCH_DEFAULT_DIRS = 0x00001000 

21 

22POINTER_ADDRESS_SPACE = 2 ** (struct.calcsize("P") * 8) 

23 

24# Set up kernel32 functions with proper types 

25kernel32 = ctypes.windll.kernel32 # type: ignore[attr-defined] 

26 

27# GetModuleHandleW 

28kernel32.GetModuleHandleW.argtypes = [ctypes.wintypes.LPCWSTR] 

29kernel32.GetModuleHandleW.restype = ctypes.wintypes.HMODULE 

30 

31# LoadLibraryExW 

32kernel32.LoadLibraryExW.argtypes = [ 

33 ctypes.wintypes.LPCWSTR, # lpLibFileName 

34 ctypes.wintypes.HANDLE, # hFile (reserved, must be NULL) 

35 ctypes.wintypes.DWORD, # dwFlags 

36] 

37kernel32.LoadLibraryExW.restype = ctypes.wintypes.HMODULE 

38 

39# GetModuleFileNameW 

40kernel32.GetModuleFileNameW.argtypes = [ 

41 ctypes.wintypes.HMODULE, # hModule 

42 ctypes.wintypes.LPWSTR, # lpFilename 

43 ctypes.wintypes.DWORD, # nSize 

44] 

45kernel32.GetModuleFileNameW.restype = ctypes.wintypes.DWORD 

46 

47 

48def ctypes_handle_to_unsigned_int(handle: ctypes.wintypes.HMODULE) -> int: 

49 """Convert ctypes HMODULE to unsigned int.""" 

50 handle_uint = int(handle) 1abc

51 if handle_uint < 0: 1abc

52 # Convert from signed to unsigned representation 

53 handle_uint += POINTER_ADDRESS_SPACE 

54 return handle_uint 1abc

55 

56 

57def add_dll_directory(dll_abs_path: str) -> None: 

58 """Add a DLL directory to the search path and update PATH environment variable. 

59 

60 Args: 

61 dll_abs_path: Absolute path to the DLL file 

62 

63 Raises: 

64 AssertionError: If the directory containing the DLL does not exist 

65 """ 

66 dirpath = os.path.dirname(dll_abs_path) 

67 assert os.path.isdir(dirpath), dll_abs_path 

68 

69 # Add the DLL directory to the native search path via the stdlib wrapper 

70 # around AddDllDirectory. This only affects the LOAD_LIBRARY_SEARCH_USER_DIRS 

71 # search; PATH is updated unconditionally below to also cover legacy 

72 # dependent-DLL resolution. The returned handle is intentionally discarded: 

73 # the directory must stay on the search path for the process lifetime, and 

74 # the handle has no finalizer, so dropping it does not remove the directory. 

75 try: 

76 os.add_dll_directory(dirpath) # type: ignore[attr-defined] 

77 except OSError as e: 

78 # Warn instead of failing silently; the PATH update below is a weaker 

79 # fallback that newer loaders may ignore. 

80 warnings.warn( 

81 f"os.add_dll_directory({dirpath!r}) failed ({e}); " 

82 "falling back to process-global PATH mutation for dependent-DLL resolution.", 

83 RuntimeWarning, 

84 stacklevel=2, 

85 ) 

86 

87 # Update PATH as a fallback for dependent DLL resolution 

88 curr_path = os.environ.get("PATH") 

89 os.environ["PATH"] = dirpath if curr_path is None else os.pathsep.join((curr_path, dirpath)) 

90 

91 

92def abs_path_for_dynamic_library(libname: str, handle: ctypes.wintypes.HMODULE) -> str: 

93 """Get the absolute path of a loaded dynamic library on Windows.""" 

94 # Create buffer for the path 

95 buffer = ctypes.create_unicode_buffer(260) # MAX_PATH 1abc

96 length = kernel32.GetModuleFileNameW(handle, buffer, len(buffer)) 1abc

97 

98 if length == 0: 1abc

99 error_code = ctypes.GetLastError() # type: ignore[attr-defined] 

100 raise RuntimeError(f"GetModuleFileNameW failed for {libname!r} (error code: {error_code})") 

101 

102 # If buffer was too small, try with larger buffer 

103 if length == len(buffer): 1abc

104 buffer = ctypes.create_unicode_buffer(32768) # Extended path length 

105 length = kernel32.GetModuleFileNameW(handle, buffer, len(buffer)) 

106 if length == 0: 

107 error_code = ctypes.GetLastError() # type: ignore[attr-defined] 

108 raise RuntimeError(f"GetModuleFileNameW failed for {libname!r} (error code: {error_code})") 

109 

110 return buffer.value 1abc

111 

112 

113def check_if_already_loaded_from_elsewhere(desc: LibDescriptor, have_abs_path: bool) -> LoadedDL | None: 

114 for dll_name in desc.windows_dlls: 1abc

115 handle = kernel32.GetModuleHandleW(dll_name) 1abc

116 if handle: 1abc

117 abs_path = abs_path_for_dynamic_library(desc.name, handle) 

118 if have_abs_path and desc.requires_add_dll_directory: 

119 # This is a side-effect if the pathfinder loads the library via 

120 # load_with_abs_path(). To make the side-effect more deterministic, 

121 # activate it even if the library was already loaded from elsewhere. 

122 add_dll_directory(abs_path) 

123 return LoadedDL(abs_path, True, ctypes_handle_to_unsigned_int(handle), "was-already-loaded-from-elsewhere") 

124 return None 1abc

125 

126 

127def load_with_system_search(desc: LibDescriptor) -> LoadedDL | None: 

128 """Try to load a DLL using the native Windows process DLL search path. 

129 

130 This calls ``LoadLibraryExW(dll_name, NULL, 0)`` directly. Under Python 

131 3.8+, CPython configures the process with 

132 ``SetDefaultDllDirectories(LOAD_LIBRARY_SEARCH_DEFAULT_DIRS)``, so this 

133 search does **not** include the system ``PATH``. Directories added via 

134 ``AddDllDirectory()`` still participate. 

135 

136 Args: 

137 desc: Descriptor for the library to load 

138 

139 Returns: 

140 A LoadedDL object if successful, None if the library cannot be loaded 

141 """ 

142 # Reverse tabulated names to achieve new -> old search order. 

143 for dll_name in reversed(desc.windows_dlls): 1abc

144 handle = kernel32.LoadLibraryExW(dll_name, None, 0) 1abc

145 if handle: 1abc

146 abs_path = abs_path_for_dynamic_library(desc.name, handle) 1abc

147 return LoadedDL(abs_path, False, ctypes_handle_to_unsigned_int(handle), "system-search") 1abc

148 

149 return None 

150 

151 

152def load_with_abs_path(desc: LibDescriptor, found_path: str, found_via: str | None = None) -> LoadedDL: 

153 """Load a dynamic library from the given path. 

154 

155 Args: 

156 desc: Descriptor for the library to load. 

157 found_path: The absolute path to the DLL file. 

158 found_via: Label indicating how the path was discovered. 

159 

160 Returns: 

161 A LoadedDL object representing the loaded library. 

162 

163 Raises: 

164 RuntimeError: If the DLL cannot be loaded. 

165 """ 

166 if desc.requires_add_dll_directory: 

167 add_dll_directory(found_path) 

168 

169 flags = WINBASE_LOAD_LIBRARY_SEARCH_DEFAULT_DIRS | WINBASE_LOAD_LIBRARY_SEARCH_DLL_LOAD_DIR 

170 handle = kernel32.LoadLibraryExW(found_path, None, flags) 

171 

172 if not handle: 

173 error_code = ctypes.GetLastError() # type: ignore[attr-defined] 

174 raise RuntimeError(f"Failed to load DLL at {found_path}: Windows error {error_code}") 

175 

176 return LoadedDL(found_path, False, ctypes_handle_to_unsigned_int(handle), found_via)