Coverage for cuda / pathfinder / _static_libs / find_bitcode_lib.py: 98.70%

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

2# SPDX-License-Identifier: Apache-2.0 

3 

4import functools (empty)

5import os (empty)

6from dataclasses import dataclass (empty)

7from typing import NoReturn, TypedDict (empty)

8 

9from cuda.pathfinder._utils.env_vars import get_cuda_home_or_path (empty)

10from cuda.pathfinder._utils.find_sub_dirs import find_sub_dirs_all_sitepackages (empty)

11from cuda.pathfinder._utils.platform_aware import IS_WINDOWS (empty)

12 

13 

14class BitcodeLibNotFoundError(RuntimeError): (empty)

15 """Raised when a bitcode library cannot be found.""" 

16 

17 

18@dataclass(frozen=True) (empty)

19class LocatedBitcodeLib: (empty)

20 """Information about a located bitcode library.""" 

21 

22 name: str 

23 abs_path: str 

24 filename: str 

25 found_via: str 

26 

27 

28class _BitcodeLibInfo(TypedDict): (empty)

29 filename: str 

30 rel_path: str 

31 site_packages_dirs: tuple[str, ...] 

32 

33 

34_SUPPORTED_BITCODE_LIBS_INFO: dict[str, _BitcodeLibInfo] = { (empty)

35 "device": { 

36 "filename": "libdevice.10.bc", 

37 "rel_path": os.path.join("nvvm", "libdevice"), 

38 "site_packages_dirs": ( 

39 "nvidia/cu13/nvvm/libdevice", 

40 "nvidia/cuda_nvcc/nvvm/libdevice", 

41 ), 

42 }, 

43} 

44 

45# Public API: just the supported library names 

46SUPPORTED_BITCODE_LIBS: tuple[str, ...] = tuple(sorted(_SUPPORTED_BITCODE_LIBS_INFO.keys())) (empty)

47 

48 

49def _no_such_file_in_dir(dir_path: str, filename: str, error_messages: list[str], attachments: list[str]) -> None: (empty)

50 error_messages.append(f"No such file: {os.path.join(dir_path, filename)}") 1 ctx1a

51 if os.path.isdir(dir_path): 1 ctx1a

52 attachments.append(f' listdir("{dir_path}"):') 1 ctx1a

53 for node in sorted(os.listdir(dir_path)): 1 ctx1a

54 attachments.append(f" {node}") 1 ctx1a

55 else: 

56 attachments.append(f' Directory does not exist: "{dir_path}"') 

57 

58 

59class _FindBitcodeLib: (empty)

60 def __init__(self, name: str) -> None: (empty)

61 if name not in _SUPPORTED_BITCODE_LIBS_INFO: # Updated reference 5 ctx1facdb

62 raise ValueError(f"Unknown bitcode library: '{name}'. Supported: {', '.join(SUPPORTED_BITCODE_LIBS)}") 1 ctx1f

63 self.name: str = name 4 ctx1acdb

64 self.config: _BitcodeLibInfo = _SUPPORTED_BITCODE_LIBS_INFO[name] # Updated reference 4 ctx1acdb

65 self.filename: str = self.config["filename"] 4 ctx1acdb

66 self.rel_path: str = self.config["rel_path"] 4 ctx1acdb

67 self.site_packages_dirs: tuple[str, ...] = self.config["site_packages_dirs"] 4 ctx1acdb

68 self.error_messages: list[str] = [] 4 ctx1acdb

69 self.attachments: list[str] = [] 4 ctx1acdb

70 

71 def try_site_packages(self) -> str | None: (empty)

72 for rel_dir in self.site_packages_dirs: 4 ctx1acdb

73 sub_dir = tuple(rel_dir.split("/")) 4 ctx1acdb

74 for abs_dir in find_sub_dirs_all_sitepackages(sub_dir): 4 ctx1acdb

75 file_path = os.path.join(abs_dir, self.filename) 1 ctx1b

76 if os.path.isfile(file_path): 1 ctx1b

77 return file_path 1 ctx1b

78 return None 4 ctx1acdb

79 

80 def try_with_conda_prefix(self) -> str | None: (empty)

81 conda_prefix = os.environ.get("CONDA_PREFIX") 4 ctx1acdb

82 if not conda_prefix: 4 ctx1acdb

83 return None 3 ctx1acd

84 

85 anchor = os.path.join(conda_prefix, "Library") if IS_WINDOWS else conda_prefix 1 ctx1b

86 file_path = os.path.join(anchor, self.rel_path, self.filename) 1 ctx1b

87 if os.path.isfile(file_path): 1 ctx1b

88 return file_path 1 ctx1b

89 return None 1 ctx1b

90 

91 def try_with_cuda_home(self) -> str | None: (empty)

92 cuda_home = get_cuda_home_or_path() 4 ctx1acdb

93 if cuda_home is None: 4 ctx1acdb

94 self.error_messages.append("CUDA_HOME/CUDA_PATH not set") 1 ctx1c

95 return None 1 ctx1c

96 

97 file_path = os.path.join(cuda_home, self.rel_path, self.filename) 3 ctx1adb

98 if os.path.isfile(file_path): 3 ctx1adb

99 return file_path 2 ctx1db

100 

101 _no_such_file_in_dir( 1 ctx1a

102 os.path.join(cuda_home, self.rel_path), 

103 self.filename, 

104 self.error_messages, 

105 self.attachments, 

106 ) 

107 return None 1 ctx1a

108 

109 def raise_not_found_error(self) -> NoReturn: (empty)

110 err = ", ".join(self.error_messages) if self.error_messages else "No search paths available" 2 ctx1ac

111 att = "\n".join(self.attachments) if self.attachments else "" 2 ctx1ac

112 raise BitcodeLibNotFoundError(f'Failure finding "{self.filename}": {err}\n{att}') 2 ctx1ac

113 

114 

115def locate_bitcode_lib(name: str) -> LocatedBitcodeLib: (empty)

116 """Locate a bitcode library by name. 

117 

118 Raises: 

119 ValueError: If ``name`` is not a supported bitcode library. 

120 BitcodeLibNotFoundError: If the bitcode library cannot be found. 

121 """ 

122 finder = _FindBitcodeLib(name) 5 ctx1facdb

123 

124 abs_path = finder.try_site_packages() 4 ctx1acdb

125 if abs_path is not None: 4 ctx1acdb

126 return LocatedBitcodeLib( 1 ctx1b

127 name=name, 

128 abs_path=abs_path, 

129 filename=finder.filename, 

130 found_via="site-packages", 

131 ) 

132 

133 abs_path = finder.try_with_conda_prefix() 4 ctx1acdb

134 if abs_path is not None: 4 ctx1acdb

135 return LocatedBitcodeLib( 1 ctx1b

136 name=name, 

137 abs_path=abs_path, 

138 filename=finder.filename, 

139 found_via="conda", 

140 ) 

141 

142 abs_path = finder.try_with_cuda_home() 4 ctx1acdb

143 if abs_path is not None: 4 ctx1acdb

144 return LocatedBitcodeLib( 2 ctx1db

145 name=name, 

146 abs_path=abs_path, 

147 filename=finder.filename, 

148 found_via="CUDA_HOME", 

149 ) 

150 

151 finder.raise_not_found_error() 2 ctx1ac

152 

153 

154@functools.cache (empty)

155def find_bitcode_lib(name: str) -> str: (empty)

156 """Find the absolute path to a bitcode library. 

157 

158 Raises: 

159 ValueError: If ``name`` is not a supported bitcode library. 

160 BitcodeLibNotFoundError: If the bitcode library cannot be found. 

161 """ 

162 return locate_bitcode_lib(name).abs_path 3 ctx1acd