Coverage for cuda/core/_memory/_copy_enums.py: 42.59%
54 statements
« prev ^ index » next coverage.py v7.16.2, created at 2026-10-16 01:23 +0000
1# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
2#
3# SPDX-License-Identifier: Apache-2.0
5from __future__ import annotations
7import dataclasses
8from collections.abc import Sequence
9from enum import StrEnum
11from cuda.core._device import Device
12from cuda.core._host import Host
13from cuda.core._utils.cuda_utils import driver
15__all__ = ["CopyOptions", "MemcpyOverlapMode", "MemcpySrcAccessOrder"]
18class MemcpySrcAccessOrder(StrEnum):
19 """Source access order hint for batched memcpy operations.
21 Maps to ``CUmemcpySrcAccessOrder``.
23 ``STREAM``
24 Source reads follow stream order. Earlier stream work may still be
25 accessing the source when the copy is enqueued.
26 ``DURING_API_CALL``
27 The driver may read the source out of stream order, but all reads
28 are complete before :func:`copy_batch` returns. No earlier stream
29 work may be accessing the source at the time of the call.
30 ``ANY``
31 The driver may read the source after the call returns. The caller
32 must keep the source unchanged until the copy completes in stream
33 order. No earlier stream work may be accessing the source.
34 """
36 STREAM = "stream"
37 DURING_API_CALL = "during_api_call"
38 ANY = "any"
41class MemcpyOverlapMode(StrEnum):
42 """Overlap mode hint for batched memcpy operations.
44 Maps to ``CUmemcpyFlags``.
46 ``DEFAULT``
47 No overlap preference; the driver uses its default scheduling.
48 ``PREFER_OVERLAP_WITH_COMPUTE``
49 Hint that the copy should preferably overlap with concurrent
50 compute work. This is advisory and may be ignored depending on
51 the platform and copy parameters.
52 """
54 DEFAULT = "default"
55 PREFER_OVERLAP_WITH_COMPUTE = "prefer_overlap_with_compute"
58@dataclasses.dataclass(frozen=True)
59class CopyOptions:
60 """Attribute bundle for a single copy within a batched memcpy.
62 Parameters
63 ----------
64 src_access_order : :class:`MemcpySrcAccessOrder` or str
65 Hint describing how the source will be accessed.
66 Default is ``"stream"`` (stream-ordered access).
67 src_location_hint : :class:`cuda.core.Device` | :class:`cuda.core.Host` | None
68 Hint for the source memory location. Honored only for managed
69 memory on devices with concurrent managed access and for
70 system-allocated pageable memory on devices with pageable memory
71 access; ignored for all other memory types. Does not prefetch
72 memory and does not set persistent memory advice.
73 ``None`` means no hint.
74 dst_location_hint : :class:`cuda.core.Device` | :class:`cuda.core.Host` | None
75 Hint for the destination memory location. Same semantics and
76 restrictions as ``src_location_hint``. ``None`` means no hint.
77 overlap_mode : :class:`MemcpyOverlapMode` or str
78 Hint requesting that the copy overlap with concurrent compute work.
79 This is advisory; it has an effect only on devices that support it.
80 Default is ``"default"``.
81 """
83 src_access_order: MemcpySrcAccessOrder | str = "stream"
84 src_location_hint: Device | Host | None = None
85 dst_location_hint: Device | Host | None = None
86 overlap_mode: MemcpyOverlapMode | str = "default"
88 def __post_init__(self):
89 # Frozen, unlike the other *Options dataclasses in cuda.core, because
90 # the batched-API contract agreed in NVIDIA/cuda-python#1775 specifies
91 # immutable per-call options:
92 # https://github.com/NVIDIA/cuda-python/pull/1775#issuecomment-4355502334
93 #
94 # Normalizing str -> StrEnum therefore has to go through
95 # object.__setattr__; a plain assignment would raise
96 # FrozenInstanceError. Done here rather than at use so that a typo
97 # fails at construction and the field always holds the enum.
98 if not isinstance(self.src_access_order, MemcpySrcAccessOrder): 1abcdefghijknlmopqrstCDRzwxuEySvFGHABTIJUKVLMNOPQ
99 try: 1abcdefghijklmqCDRzuEvABL
100 object.__setattr__( 1abcdefghijklmqCDRzuEvABL
101 self,
102 "src_access_order",
103 MemcpySrcAccessOrder(self.src_access_order),
104 )
105 except (ValueError, TypeError) as exc: 1R
106 raise ValueError(f"invalid src_access_order: {self.src_access_order!r}") from exc 1R
107 if not isinstance(self.overlap_mode, MemcpyOverlapMode): 1abcdefghijknlmopqrstCDzwxuEySvFGHABTIJUKVLMNOPQ
108 try: 1abcdefghijknlmoprstCDzwxuEySvFGHABTIJUKVMNOPQ
109 object.__setattr__( 1abcdefghijknlmoprstCDzwxuEySvFGHABTIJUKVMNOPQ
110 self,
111 "overlap_mode",
112 MemcpyOverlapMode(self.overlap_mode),
113 )
114 except (ValueError, TypeError) as exc: 1z
115 raise ValueError(f"invalid overlap_mode: {self.overlap_mode!r}") from exc 1z
117 def _to_driver_enum(self) -> int:
118 """Return the driver CUmemcpySrcAccessOrder value."""
119 return _SRC_ACCESS_ORDER_TO_DRIVER[MemcpySrcAccessOrder(self.src_access_order)] 1abcdefghijknlmopqrstFGHABIJKLMNOPQ
121 def _to_driver_flags(self) -> int:
122 """Return the driver CUmemcpyFlags value."""
123 return _OVERLAP_MODE_TO_DRIVER[MemcpyOverlapMode(self.overlap_mode)] 1abcdefghijknlmopqrstFGHABIJKLMNOPQ
126# CUDA 12.8 added CUmemcpySrcAccessOrder and CUmemcpyFlags. Every
127# cuda-bindings that cuda.core accepts has them. Keyed by ``str`` (StrEnum
128# members are ``str``) and valued as ``int``.
129_src_order = driver.CUmemcpySrcAccessOrder
130_flags = driver.CUmemcpyFlags
131_SRC_ACCESS_ORDER_TO_DRIVER: dict[str, int] = {
132 MemcpySrcAccessOrder.STREAM: int(_src_order.CU_MEMCPY_SRC_ACCESS_ORDER_STREAM),
133 MemcpySrcAccessOrder.DURING_API_CALL: int(_src_order.CU_MEMCPY_SRC_ACCESS_ORDER_DURING_API_CALL),
134 MemcpySrcAccessOrder.ANY: int(_src_order.CU_MEMCPY_SRC_ACCESS_ORDER_ANY),
135}
136_OVERLAP_MODE_TO_DRIVER: dict[str, int] = {
137 MemcpyOverlapMode.DEFAULT: int(_flags.CU_MEMCPY_FLAG_DEFAULT),
138 MemcpyOverlapMode.PREFER_OVERLAP_WITH_COMPUTE: int(_flags.CU_MEMCPY_FLAG_PREFER_OVERLAP_WITH_COMPUTE),
139}
140del _src_order, _flags
143def _reject_unsupported_during_api_call(
144 src_access_order: MemcpySrcAccessOrder, requirement: str, *, index: int | None = None
145) -> None:
146 """Raise if ``src_access_order`` is DURING_API_CALL but the native attributes
147 path (``cuMemcpyWithAttributesAsync`` / ``cuMemcpyBatchAsync``) is unavailable.
149 STREAM and ANY never promise access sooner than stream order, so a plain
150 ``cuMemcpyAsync`` fallback satisfies them; DURING_API_CALL specifically
151 promises all source reads complete before the call returns, which
152 ``cuMemcpyAsync`` cannot provide (it reads the source in stream order
153 only). Silently downgrading that guarantee would let a caller reuse or
154 overwrite the source buffer before the real, stream-ordered read
155 happens: a silent data race, not a missed optimization. ``requirement``
156 names what the native path needs and why it is unavailable here;
157 ``index`` identifies the offending copy within a batch.
159 Internal, but deliberately importable: shared between the per-buffer and
160 batched fallback paths so both raise identically, and directly testable
161 without needing an actual old driver/bindings install.
162 """
163 if src_access_order != MemcpySrcAccessOrder.DURING_API_CALL: 1WXYZ0
164 return 1Z0
165 where = f" at index {index}" if index is not None else "" 1WXY
166 raise RuntimeError( 1WXY
167 f"src_access_order=DURING_API_CALL{where} requires {requirement}. A "
168 "plain cuMemcpyAsync fallback reads the source in stream order only, "
169 "which would silently violate the guarantee that all source reads "
170 "complete before the call returns, letting the caller reuse the "
171 "source buffer before the real (stream-ordered) read happens. Use "
172 "src_access_order=STREAM or ANY, or omit options, if that works for "
173 "your use case."
174 )
177def _attr_run_starts(attrs: Sequence[CopyOptions]) -> list[int]:
178 """Return the start index of each maximal run of equal attributes.
180 This mirrors the ``attrsIdxs`` indirection that ``cuMemcpyBatchAsync``
181 expects: ``attrs[k]`` applies to the copies in
182 ``[starts[k], starts[k + 1])``. Collapsing equal neighbours means a
183 broadcast attribute is passed to the driver once (``numAttrs == 1``)
184 rather than repeated per copy.
185 """
186 starts: list[int] = [] 1abcdefghijknlmopqrstwxuyv
187 prev: CopyOptions | None = None 1abcdefghijknlmopqrstwxuyv
188 for i, attr in enumerate(attrs): 1abcdefghijknlmopqrstwxuyv
189 if i == 0 or attr != prev: 1abcdefghijknlmopqrstwxuyv
190 starts.append(i) 1abcdefghijknlmopqrstwxuyv
191 prev = attr 1abcdefghijknlmopqrstwxuyv
192 return starts 1abcdefghijknlmopqrstwxuyv