Coverage for cuda/core/_memory/_copy_enums.py: 93.85%

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

2# 

3# SPDX-License-Identifier: Apache-2.0 

4 

5from __future__ import annotations (empty)

6 

7import dataclasses (empty)

8from collections.abc import Sequence (empty)

9 

10from cuda.core._device import Device (empty)

11from cuda.core._host import Host (empty)

12from cuda.core._utils.cuda_utils import driver (empty)

13from cuda.core._utils.pycompat import StrEnum (empty)

14from cuda.core._utils.version import binding_version (empty)

15 

16__all__ = ["CopyOptions", "MemcpyOverlapMode", "MemcpySrcAccessOrder"] (empty)

17 

18 

19class MemcpySrcAccessOrder(StrEnum): (empty)

20 """Source access order hint for batched memcpy operations. 

21 

22 Maps to ``CUmemcpySrcAccessOrder``. 

23 

24 ``STREAM`` 

25 Source reads follow stream order. Earlier stream work may still be 

26 accessing the source when the copy is enqueued. 

27 ``DURING_API_CALL`` 

28 The driver may read the source out of stream order, but all reads 

29 are complete before :func:`copy_batch` returns. No earlier stream 

30 work may be accessing the source at the time of the call. 

31 ``ANY`` 

32 The driver may read the source after the call returns. The caller 

33 must keep the source unchanged until the copy completes in stream 

34 order. No earlier stream work may be accessing the source. 

35 """ 

36 

37 STREAM = "stream" (empty)

38 DURING_API_CALL = "during_api_call" (empty)

39 ANY = "any" (empty)

40 

41 

42class MemcpyOverlapMode(StrEnum): (empty)

43 """Overlap mode hint for batched memcpy operations. 

44 

45 Maps to ``CUmemcpyFlags``. 

46 

47 ``DEFAULT`` 

48 No overlap preference; the driver uses its default scheduling. 

49 ``PREFER_OVERLAP_WITH_COMPUTE`` 

50 Hint that the copy should preferably overlap with concurrent 

51 compute work. This is advisory and may be ignored depending on 

52 the platform and copy parameters. 

53 """ 

54 

55 DEFAULT = "default" (empty)

56 PREFER_OVERLAP_WITH_COMPUTE = "prefer_overlap_with_compute" (empty)

57 

58 

59@dataclasses.dataclass(frozen=True) (empty)

60class CopyOptions: (empty)

61 """Attribute bundle for a single copy within a batched memcpy. 

62 

63 Parameters 

64 ---------- 

65 src_access_order : :class:`MemcpySrcAccessOrder` or str 

66 Hint describing how the source will be accessed. 

67 Default is ``"stream"`` (stream-ordered access). 

68 src_location_hint : :class:`cuda.core.Device` | :class:`cuda.core.Host` | None 

69 Hint for the source memory location. Honored only for managed 

70 memory on devices with concurrent managed access and for 

71 system-allocated pageable memory on devices with pageable memory 

72 access; ignored for all other memory types. Does not prefetch 

73 memory and does not set persistent memory advice. 

74 ``None`` means no hint. 

75 dst_location_hint : :class:`cuda.core.Device` | :class:`cuda.core.Host` | None 

76 Hint for the destination memory location. Same semantics and 

77 restrictions as ``src_location_hint``. ``None`` means no hint. 

78 overlap_mode : :class:`MemcpyOverlapMode` or str 

79 Hint requesting that the copy overlap with concurrent compute work. 

80 This is advisory; it has an effect only on devices that support it. 

81 Default is ``"default"``. 

82 """ 

83 

84 src_access_order: MemcpySrcAccessOrder | str = "stream" (empty)

85 src_location_hint: Device | Host | None = None (empty)

86 dst_location_hint: Device | Host | None = None (empty)

87 overlap_mode: MemcpyOverlapMode | str = "default" (empty)

88 

89 def __post_init__(self): (empty)

90 # Frozen, unlike the other *Options dataclasses in cuda.core, because 

91 # the batched-API contract agreed in NVIDIA/cuda-python#1775 specifies 

92 # immutable per-call options: 

93 # https://github.com/NVIDIA/cuda-python/pull/1775#issuecomment-4355502334 

94 # 

95 # Normalizing str -> StrEnum therefore has to go through 

96 # object.__setattr__; a plain assignment would raise 

97 # FrozenInstanceError. Done here rather than at use so that a typo 

98 # fails at construction and the field always holds the enum. 

99 if not isinstance(self.src_access_order, MemcpySrcAccessOrder): 48 ctx1bcdefghijklomnpqrstuPQSCxyvRzTwDEFABUGHVIWJKLMNO

100 try: 24 ctx1bcdefghijklmnrPQSCvRwABJ

101 object.__setattr__( 24 ctx1bcdefghijklmnrPQSCvRwABJ

102 self, 

103 "src_access_order", 

104 MemcpySrcAccessOrder(self.src_access_order), 

105 ) 

106 except (ValueError, TypeError) as exc: 1 ctx1S

107 raise ValueError(f"invalid src_access_order: {self.src_access_order!r}") from exc 1 ctx1S

108 if not isinstance(self.overlap_mode, MemcpyOverlapMode): 47 ctx1bcdefghijklomnpqrstuPQCxyvRzTwDEFABUGHVIWJKLMNO

109 try: 45 ctx1bcdefghijklomnpqstuPQCxyvRzTwDEFABUGHVIWKLMNO

110 object.__setattr__( 45 ctx1bcdefghijklomnpqstuPQCxyvRzTwDEFABUGHVIWKLMNO

111 self, 

112 "overlap_mode", 

113 MemcpyOverlapMode(self.overlap_mode), 

114 ) 

115 except (ValueError, TypeError) as exc: 1 ctx1C

116 raise ValueError(f"invalid overlap_mode: {self.overlap_mode!r}") from exc 1 ctx1C

117 

118 def _to_driver_enum(self) -> int: (empty)

119 """Return the driver CUmemcpySrcAccessOrder value.""" 

120 if not _SRC_ACCESS_ORDER_TO_DRIVER: 34 ctx1bcdefghijklomnpqrstuDEFABGHIJKLMNO

121 raise NotImplementedError(_CUDA13_REQUIRED) 

122 return _SRC_ACCESS_ORDER_TO_DRIVER[MemcpySrcAccessOrder(self.src_access_order)] 34 ctx1bcdefghijklomnpqrstuDEFABGHIJKLMNO

123 

124 def _to_driver_flags(self) -> int: (empty)

125 """Return the driver CUmemcpyFlags value.""" 

126 if not _OVERLAP_MODE_TO_DRIVER: 34 ctx1bcdefghijklomnpqrstuDEFABGHIJKLMNO

127 raise NotImplementedError(_CUDA13_REQUIRED) 

128 return _OVERLAP_MODE_TO_DRIVER[MemcpyOverlapMode(self.overlap_mode)] 34 ctx1bcdefghijklomnpqrstuDEFABGHIJKLMNO

129 

130 

131_CUDA13_REQUIRED = "copy attributes require cuda.bindings 13.0 or newer" (empty)

132 

133# CUmemcpySrcAccessOrder and CUmemcpyFlags are exposed by cuda.bindings 13.0+, 

134# so these maps are empty when it is older. Nothing reaches them there: 

135# copy_batch refuses non-default CopyOptions when the batched entry point is 

136# unavailable. 

137# 

138# Keyed by ``str``: under ``python_version = "3.10"`` mypy resolves StrEnum to 

139# the unstubbed backports shim and so infers the members as plain ``str``. 

140# StrEnum members are ``str`` instances, so this holds on every version. The 

141# values are wrapped in ``int()`` because the driver enums are untyped. 

142_SRC_ACCESS_ORDER_TO_DRIVER: dict[str, int] (empty)

143_OVERLAP_MODE_TO_DRIVER: dict[str, int] (empty)

144 

145if binding_version() >= (13, 0, 0): (empty)

146 _src_order = driver.CUmemcpySrcAccessOrder (empty)

147 _flags = driver.CUmemcpyFlags (empty)

148 _SRC_ACCESS_ORDER_TO_DRIVER = { (empty)

149 MemcpySrcAccessOrder.STREAM: int(_src_order.CU_MEMCPY_SRC_ACCESS_ORDER_STREAM), 

150 MemcpySrcAccessOrder.DURING_API_CALL: int(_src_order.CU_MEMCPY_SRC_ACCESS_ORDER_DURING_API_CALL), 

151 MemcpySrcAccessOrder.ANY: int(_src_order.CU_MEMCPY_SRC_ACCESS_ORDER_ANY), 

152 } 

153 _OVERLAP_MODE_TO_DRIVER = { (empty)

154 MemcpyOverlapMode.DEFAULT: int(_flags.CU_MEMCPY_FLAG_DEFAULT), 

155 MemcpyOverlapMode.PREFER_OVERLAP_WITH_COMPUTE: int(_flags.CU_MEMCPY_FLAG_PREFER_OVERLAP_WITH_COMPUTE), 

156 } 

157 del _src_order, _flags (empty)

158else: 

159 _SRC_ACCESS_ORDER_TO_DRIVER = {} 

160 _OVERLAP_MODE_TO_DRIVER = {} 

161 

162 

163def _reject_unsupported_during_api_call( (empty)

164 src_access_order: MemcpySrcAccessOrder, requirement: str, *, index: int | None = None 

165) -> None: 

166 """Raise if ``src_access_order`` is DURING_API_CALL but the native attributes 

167 path (``cuMemcpyWithAttributesAsync`` / ``cuMemcpyBatchAsync``) is unavailable. 

168 

169 STREAM and ANY never promise access sooner than stream order, so a plain 

170 ``cuMemcpyAsync`` fallback satisfies them; DURING_API_CALL specifically 

171 promises all source reads complete before the call returns, which 

172 ``cuMemcpyAsync`` cannot provide (it reads the source in stream order 

173 only). Silently downgrading that guarantee would let a caller reuse or 

174 overwrite the source buffer before the real, stream-ordered read 

175 happens: a silent data race, not a missed optimization. ``requirement`` 

176 names what the native path needs and why it is unavailable here; 

177 ``index`` identifies the offending copy within a batch. 

178 

179 Internal, but deliberately importable: shared between the per-buffer and 

180 batched fallback paths so both raise identically, and directly testable 

181 without needing an actual old driver/bindings install. 

182 """ 

183 if src_access_order != MemcpySrcAccessOrder.DURING_API_CALL: 5 ctx1XYZ01

184 return 2 ctx101

185 where = f" at index {index}" if index is not None else "" 3 ctx1XYZ

186 raise RuntimeError( 3 ctx1XYZ

187 f"src_access_order=DURING_API_CALL{where} requires {requirement}. A " 

188 "plain cuMemcpyAsync fallback reads the source in stream order only, " 

189 "which would silently violate the guarantee that all source reads " 

190 "complete before the call returns, letting the caller reuse the " 

191 "source buffer before the real (stream-ordered) read happens. Use " 

192 "src_access_order=STREAM or ANY, or omit options, if that works for " 

193 "your use case." 

194 ) 

195 

196 

197def _attr_run_starts(attrs: Sequence[CopyOptions]) -> list[int]: (empty)

198 """Return the start index of each maximal run of equal attributes. 

199 

200 This mirrors the ``attrsIdxs`` indirection that ``cuMemcpyBatchAsync`` 

201 expects: ``attrs[k]`` applies to the copies in 

202 ``[starts[k], starts[k + 1])``. Collapsing equal neighbours means a 

203 broadcast attribute is passed to the driver once (``numAttrs == 1``) 

204 rather than repeated per copy. 

205 """ 

206 starts: list[int] = [] 25 ctx1bcdefghijklomnpqrstuxyvzw

207 prev: CopyOptions | None = None 25 ctx1bcdefghijklomnpqrstuxyvzw

208 for i, attr in enumerate(attrs): 25 ctx1bcdefghijklomnpqrstuxyvzw

209 if i == 0 or attr != prev: 25 ctx1bcdefghijklomnpqrstuxyvzw

210 starts.append(i) 25 ctx1bcdefghijklomnpqrstuxyvzw

211 prev = attr 25 ctx1bcdefghijklomnpqrstuxyvzw

212 return starts 25 ctx1bcdefghijklomnpqrstuxyvzw