Coverage for cuda/core/typing.py: 98.46%

65 statements  

« prev     ^ index     » next       coverage.py v7.14.1, created at 2026-06-13 01:38 +0000

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

2# 

3# SPDX-License-Identifier: Apache-2.0 

4 

5"""Public type aliases, protocols, and enumerations used in cuda.core API signatures.""" 

6 

7import sys 

8from typing import TYPE_CHECKING 

9from typing import Literal as _Literal 

10from typing import TypeAlias as _TypeAlias 

11 

12if TYPE_CHECKING: 

13 # `backports.strenum` ships no type stubs and typeshed conditionally gates 

14 # `enum.StrEnum` behind `sys.version_info >= (3, 11)`. Declaring a minimal 

15 # local shape here (mirroring typeshed's 3.11 StrEnum) lets mypy at 

16 # `python_version = "3.10"` infer subclass members as `Literal[Foo.MEMBER]` 

17 # rather than bare `str`. 

18 from enum import Enum 

19 

20 class StrEnum(str, Enum): 

21 _value_: str 

22 

23 

24if not TYPE_CHECKING: 

25 if sys.version_info >= (3, 11): 

26 from enum import StrEnum 

27 else: 

28 from backports.strenum import StrEnum 

29 

30from cuda.core._context import DeviceResourcesType 

31from cuda.core._stream import IsStreamType 

32from cuda.core._utils.cuda_utils import driver 

33 

34__all__ = [ 

35 "CompilerBackendType", 

36 "DevicePointerType", 

37 "DeviceResourcesType", 

38 "GraphConditionalType", 

39 "GraphMemoryType", 

40 "IsStreamType", 

41 "ManagedMemoryLocationType", 

42 "ObjectCodeFormatType", 

43 "PCHStatusType", 

44 "ProcessStateType", 

45 "SourceCodeType", 

46 "VirtualMemoryAccessType", 

47 "VirtualMemoryAllocationType", 

48 "VirtualMemoryGranularityType", 

49 "VirtualMemoryHandleType", 

50 "VirtualMemoryLocationType", 

51] 

52 

53 

54# A type union of :obj:`~driver.CUdeviceptr`, `int` and `None` for hinting 

55# :attr:`Buffer.handle`. 

56DevicePointerType: _TypeAlias = driver.CUdeviceptr | int | None 

57 

58 

59ProcessStateType = _Literal["running", "locked", "checkpointed", "failed"] 

60 

61 

62class SourceCodeType(StrEnum): 

63 """Source language passed to :class:`~cuda.core.Program`. 

64 

65 * ``CXX`` — CUDA C++ source. 

66 * ``PTX`` — PTX assembly text. 

67 * ``NVVM`` — NVVM IR (LLVM bitcode). 

68 """ 

69 

70 CXX = "c++" 

71 PTX = "ptx" 

72 NVVM = "nvvm" 

73 

74 

75class ObjectCodeFormatType(StrEnum): 

76 """Output format for :meth:`~cuda.core.Program.compile`, :meth:`~cuda.core.Linker.link`, and :meth:`~cuda.core.Program.as_bytes`. 

77 

78 * ``PTX`` — PTX assembly text. 

79 * ``CUBIN`` — device-native CUDA binary. 

80 * ``LTOIR`` — LTO (link-time optimization) IR for later linking. 

81 * ``FATBIN`` — fat binary bundling multiple device images. 

82 * ``OBJECT`` — relocatable device object. 

83 * ``LIBRARY`` — device code library. 

84 """ 

85 

86 PTX = "ptx" 

87 CUBIN = "cubin" 

88 LTOIR = "ltoir" 

89 FATBIN = "fatbin" 

90 OBJECT = "object" 

91 LIBRARY = "library" 

92 

93 

94class CompilerBackendType(StrEnum): 

95 """Compiler backend inferred from the program's code type and exposed on :attr:`~cuda.core.Program.backend`. 

96 

97 * ``NVRTC`` — NVIDIA Runtime Compilation. 

98 * ``NVVM`` — NVVM LLVM backend. 

99 * ``NVJITLINK`` — nvJitLink device-side linker. 

100 * ``DRIVER`` — CUDA driver PTX JIT compiler. 

101 """ 

102 

103 NVRTC = "NVRTC" 

104 NVVM = "NVVM" 

105 NVJITLINK = "nvJitLink" 

106 DRIVER = "driver" 

107 

108 

109class PCHStatusType(StrEnum): 

110 """Precompiled-header (PCH) outcome reported by :meth:`~cuda.core.Program.compile`. 

111 

112 * ``CREATED`` — PCH was successfully written. 

113 * ``NOT_ATTEMPTED`` — PCH creation was skipped (backend does not support it or the option was not requested). 

114 * ``FAILED`` — PCH creation was attempted but failed. 

115 """ 

116 

117 CREATED = "created" 

118 NOT_ATTEMPTED = "not_attempted" 

119 FAILED = "failed" 

120 

121 

122class GraphConditionalType(StrEnum): 

123 """Conditional node flavor for :class:`~cuda.core.graph.GraphBuilder`. 

124 

125 * ``IF`` — body graph executes at most once based on a condition. 

126 * ``WHILE`` — body graph loops while the condition is true. 

127 * ``SWITCH`` — selects one child graph by an integer index. 

128 """ 

129 

130 IF = "if" 

131 WHILE = "while" 

132 SWITCH = "switch" 

133 

134 

135class GraphMemoryType(StrEnum): 

136 """Memory space for a graph memory-allocation or free node. 

137 

138 * ``DEVICE`` — GPU device memory. 

139 * ``HOST`` — pinned host memory. 

140 * ``MANAGED`` — CUDA managed (unified) memory. 

141 """ 

142 

143 DEVICE = "device" 

144 HOST = "host" 

145 MANAGED = "managed" 

146 

147 

148class ManagedMemoryLocationType(StrEnum): 

149 """Destination type for managed-memory prefetch and advise operations. 

150 

151 * ``DEVICE`` — target a GPU device. 

152 * ``HOST`` — target the CPU host (any NUMA node). 

153 * ``HOST_NUMA`` — target a specific host NUMA node. 

154 """ 

155 

156 DEVICE = "device" 

157 HOST = "host" 

158 HOST_NUMA = "host_numa" 

159 

160 

161class VirtualMemoryHandleType(StrEnum): 

162 """OS handle type for exporting virtual memory allocations across processes. 

163 

164 * ``POSIX_FD`` — POSIX file descriptor (Linux). 

165 * ``WIN32_KMT`` — Win32 kernel-mode handle (Windows). 

166 * ``FABRIC`` — NVLink/NVSwitch fabric handle for multi-node topologies. 

167 """ 

168 

169 POSIX_FD = "posix_fd" 

170 WIN32_KMT = "win32_kmt" 

171 FABRIC = "fabric" 

172 

173 

174class VirtualMemoryLocationType(StrEnum): 

175 """Physical backing location for a virtual memory allocation. 

176 

177 * ``DEVICE`` — GPU device memory. 

178 * ``HOST`` — pinned host memory. 

179 * ``HOST_NUMA`` — host memory pinned to a specific NUMA node. 

180 * ``HOST_NUMA_CURRENT`` — host memory on the calling thread's NUMA node. 

181 """ 

182 

183 DEVICE = "device" 

184 HOST = "host" 

185 HOST_NUMA = "host_numa" 

186 HOST_NUMA_CURRENT = "host_numa_current" 

187 

188 

189class VirtualMemoryGranularityType(StrEnum): 

190 """Granularity query type for virtual memory allocations. 

191 

192 * ``MINIMUM`` — smallest allocation size supported by the device. 

193 * ``RECOMMENDED`` — granularity that yields best performance on the device. 

194 """ 

195 

196 MINIMUM = "minimum" 

197 RECOMMENDED = "recommended" 

198 

199 

200class VirtualMemoryAccessType(StrEnum): 

201 """Access permissions for a virtual memory mapping. 

202 

203 * ``READ_WRITE`` — both read and write access. 

204 * ``READ`` — read-only access. 

205 """ 

206 

207 READ_WRITE = "rw" 

208 READ = "r" 

209 

210 

211class VirtualMemoryAllocationType(StrEnum): 

212 """Physical memory type for a virtual memory backing allocation. 

213 

214 * ``PINNED`` — pinned/non-migratable physical allocation (placement via :class:`VirtualMemoryLocationType`). 

215 * ``MANAGED`` — CUDA managed (unified) memory (CUDA 13+ only). 

216 """ 

217 

218 PINNED = "pinned" 

219 MANAGED = "managed" 

220 

221 

222del StrEnum