"""Annotate CUDA graph kernel nodes during capture. During CUDA graph capture, ``mark_kernels`` records the current capture frontier and the direct dependents already attached to that frontier. On scope exit it walks only the newly added dependent edges to find the nodes created within the scope. Each kernel, memcpy, or memset node found is annotated by its ``toolsId`` so it can later be matched to profiler trace events. ``mark_kernels`` now snapshots capture state from whatever stream is current on scope entry, so that stream must already be participating in the capture. ``mark_stream`` handles this by starting ``mark_kernels`` before switching to the target stream. The annotations can be pickled and later merged into a Chrome profiler trace using ``torch.cuda._annotate_cuda_graph_trace``. Requires ``cuda.bindings`` package and a CUDA driver that supports ``cudaGraphNodeGetToolsId`` (CUDA >= 13.1 or appropriate cuda-compat). When unavailable, ``mark_kernels`` silently becomes a no-op. Annotation recording is enabled per capture via ``torch.cuda.graph``'s ``enable_annotations`` argument, which also resolves and remaps the recorded annotations to the exec graph automatically on context exit. Usage during capture:: from torch.cuda._graph_annotations import mark_kernels, get_kernel_annotations with torch.cuda.graph(graph, enable_annotations=True): with mark_kernels("phase_A"): y = workload_a(x) with mark_kernels("phase_B"): z = workload_b(y) annotations = get_kernel_annotations() When you need to drive this outside the context manager's automatic path, ``resolve_and_remap(graph)`` is shorthand for ``resolve_pending_annotations()`` followed by ``remap_to_exec_graph(graph)``; call those directly for finer control (e.g. resolving once before remapping several graphs). """ from __future__ import annotations import importlib.metadata import threading import warnings from collections import defaultdict from contextlib import contextmanager from logging import getLogger from typing import Any, NamedTuple, TYPE_CHECKING, TypeAlias if TYPE_CHECKING: from collections.abc import Iterable, Mapping import torch from torch.cuda._utils import ( _check_cuda_bindings, _check_cuda_bindings_driver, _HAS_CUDA_BINDINGS, ) try: from cuda.bindings import ( # pyrefly: ignore[missing-import] driver as _cuda_driver, runtime as _cuda_runtime, ) except ImportError: _cuda_driver = None # type: ignore[assignment] _cuda_runtime = None # type: ignore[assignment] logger = getLogger(__name__) _CaptureState: TypeAlias = tuple[Any, list[Any]] _ExistingDirectDependents: TypeAlias = dict[int, set[int]] # Tri-state: None = not probed, True = available, False = unavailable. # Deferred to first use to avoid premature CUDA initialization. _tools_id_available: bool | None = None # Whether annotation recording is active. Scoped to a ``torch.cuda.graph`` # capture: set by the context manager on entry (from its ``enable_annotations`` # argument) and cleared on exit. When False, mark_kernels/mark_stream and the # capture-id stamp are no-ops. Module-level (not thread-local) because capture # can span threads (e.g. autograd). _annotations_enabled: bool = False # Id of the top-level capture graph of the active capture. A conditional node's body is # captured into a separate cudaGraph_t (see CUDAGraph::begin_capture_to_conditional_node), # and its node ids live in that graph's id space, which remap_to_exec_graph never rekeys -- # so mark_kernels compares against this to detect a scope inside a body. _capture_root_graph_id: int | None = None def _set_annotations_enabled(enabled: bool) -> None: """Set whether annotation recording is active. Used by ``torch.cuda.graph`` to scope annotations to a capture; not a public API.""" global _annotations_enabled, _capture_root_graph_id _annotations_enabled = enabled if not enabled: _capture_root_graph_id = None def _graph_id(graph: Any) -> int: """The driver's unique id for a ``cudaGraph_t``. Identifies a graph across its whole lifetime, unlike the handle itself: that is a pointer, so a graph created after another is destroyed can reuse the value and compare equal to it.""" return _check_cuda_bindings( _cuda_runtime.cudaGraphGetId(graph) # pyrefly: ignore[missing-attribute] ) def maybe_stamp_capture_root(stream: Any) -> None: """Record the top-level capture graph of the capture just begun on ``stream``. Called by ``torch.cuda.graph`` right after ``capture_begin``, while ``stream`` is still capturing into the top-level graph. Not a public API.""" global _capture_root_graph_id _capture_root_graph_id = None if not _annotations_enabled or _is_tools_id_unavailable(): return stream_handle = _cuda_runtime.cudaStream_t( # pyrefly: ignore[missing-attribute] init_value=stream.cuda_stream ) state = _get_capture_state(stream_handle) if state is not None: _capture_root_graph_id = _graph_id(state[0]) def _probe_tools_id() -> bool: """Probe whether cudaGraphNodeGetToolsId is supported by the driver. Calls with a null node: cudaErrorInvalidValue means the API exists in the driver (good), cudaErrorCallRequiresNewerDriver means it does not (bad). """ if not hasattr(_cuda_runtime, "cudaGraphNodeGetToolsId"): # API is missing from cuda-bindings - likely version too old cuda_bindings_version = importlib.metadata.version("cuda-bindings") logger.warning( "cudaGraphNodeGetToolsId API not found in cuda-bindings. " "Current version: %s, required: >= 13.1.0. " "CUDA graph kernel annotations will be disabled. " "To enable annotations, upgrade cuda-bindings: " "pip install --upgrade cuda-bindings", cuda_bindings_version, ) return False err, *_ = _cuda_runtime.cudaGraphNodeGetToolsId( 0 ) # pyrefly: ignore[missing-attribute] if ( err == _cuda_runtime.cudaError_t.cudaErrorCallRequiresNewerDriver # pyrefly: ignore[missing-attribute] ): logger.info( "cudaGraphNodeGetToolsId requires a newer driver " "(missing cuda-compat?); " "CUDA graph kernel annotations will be disabled" ) return False return True def _is_tools_id_unavailable() -> bool: """Return True if cudaGraphNodeGetToolsId is not usable.""" global _tools_id_available if not _HAS_CUDA_BINDINGS: return True if _tools_id_available is not None: return not _tools_id_available _tools_id_available = _probe_tools_id() return not _tools_id_available def is_available() -> bool: r"""is_available() -> bool Return whether CUDA graph annotation recording is supported. Requires a CUDA device, the ``cuda-bindings`` package, and a driver that supports ``cudaGraphNodeGetToolsId`` (CUDA >= 13.1 or an equivalent cuda-compat package). When this returns ``False``, :func:`mark_kernels` is a silent no-op and no annotations are recorded. The first call may probe the CUDA driver; the result is cached. """ return torch.cuda.is_available() and not _is_tools_id_unavailable() def _get_capture_state(stream: Any) -> _CaptureState | None: """Return ``(graph, frontier)`` for an active capture, else ``None``.""" status, _id, graph, _deps, _edge_data, _num_deps = _check_cuda_bindings( _cuda_runtime.cudaStreamGetCaptureInfo( # pyrefly: ignore[missing-attribute] stream ) ) if ( status != _cuda_runtime.cudaStreamCaptureStatus.cudaStreamCaptureStatusActive # pyrefly: ignore[missing-attribute] ): return None return graph, list(_deps[:_num_deps]) def _get_root_nodes(graph: Any) -> list[Any]: """Return the current root nodes for the graph.""" _, num_roots = _check_cuda_bindings( _cuda_runtime.cudaGraphGetRootNodes( # pyrefly: ignore[missing-attribute] graph ) ) if num_roots == 0: return [] roots, num_roots = _check_cuda_bindings( _cuda_runtime.cudaGraphGetRootNodes( # pyrefly: ignore[missing-attribute] graph, pNumRootNodes=num_roots ) ) return list(roots[:num_roots]) def _get_dependent_nodes(node: Any) -> list[Any]: """Return the direct dependents of a graph node.""" _, _, num_dependents = _check_cuda_bindings( _cuda_runtime.cudaGraphNodeGetDependentNodes( # pyrefly: ignore[missing-attribute] node ) ) if num_dependents == 0: return [] dependents, _edge_data, num_dependents = _check_cuda_bindings( _cuda_runtime.cudaGraphNodeGetDependentNodes( # pyrefly: ignore[missing-attribute] node, pNumDependentNodes=num_dependents ) ) return list(dependents[:num_dependents]) def _get_node_type(node: Any) -> Any: """Return graph node type without tripping runtime bugs on newer node kinds. The runtime ``cudaGraphNodeGetType`` can return ``cudaErrorUnknown`` (999) for valid nodes whose driver type is ``CU_GRAPH_NODE_TYPE_BATCH_MEM_OP``. Query via the driver API instead. """ return _check_cuda_bindings_driver( _cuda_driver.cuGraphNodeGetType(node) # pyrefly: ignore[missing-attribute] ) def _collect_descendants( start_nodes: list[Any], *, existing_direct_dependents: _ExistingDirectDependents | None = None, include_start_nodes: bool = False, ) -> dict[int, Any]: """Walk dependent edges starting at ``start_nodes``. ``existing_direct_dependents`` maps each node in ``start_nodes`` to the direct dependent node keys that were already present at scope entry. Those edges are skipped so the traversal only follows nodes added after scope entry. """ existing_direct_dependents = existing_direct_dependents or {} seen = {int(node) for node in start_nodes} descendants: dict[int, Any] = {} stack = list(start_nodes) if include_start_nodes: for node in start_nodes: descendants[int(node)] = node while stack: node = stack.pop() old_dependents = existing_direct_dependents.get(int(node), set()) for dependent in _get_dependent_nodes(node): dependent_key = int(dependent) if dependent_key in old_dependents or dependent_key in seen: continue seen.add(dependent_key) descendants[dependent_key] = dependent stack.append(dependent) return descendants # toolsId -> list of annotation objects. _kernel_annotations: defaultdict[int, list[Any]] = defaultdict(list) # Node types we annotate (kernels, memcpys, memsets, batch mem ops, event # record/wait nodes, and host nodes), as driver node-type enums. Event and host # nodes are stream-ordered but do not occupy the device; annotating them lets the # profiler place their spans on the intended stream lane. Initialized lazily to # avoid touching cuda.bindings at import time. _ANNOTATABLE_TYPES: set[Any] | None = None def _get_annotatable_types() -> set[Any]: global _ANNOTATABLE_TYPES if _ANNOTATABLE_TYPES is None: node_types = _cuda_driver.CUgraphNodeType # pyrefly: ignore[missing-attribute] _ANNOTATABLE_TYPES = { node_types.CU_GRAPH_NODE_TYPE_KERNEL, node_types.CU_GRAPH_NODE_TYPE_MEMCPY, node_types.CU_GRAPH_NODE_TYPE_MEMSET, node_types.CU_GRAPH_NODE_TYPE_BATCH_MEM_OP, node_types.CU_GRAPH_NODE_TYPE_EVENT_RECORD, node_types.CU_GRAPH_NODE_TYPE_WAIT_EVENT, node_types.CU_GRAPH_NODE_TYPE_HOST, } return _ANNOTATABLE_TYPES # Node types whose work lives in a separate cudaGraph_t (child graphs, conditional # bodies). The dependent-edge walk does not descend into such a node, and the nodes # inside are numbered in the body graph's id space, which remap_to_exec_graph never # rekeys -- so annotations there would be silently lost. See mark_kernels. Initialized # lazily, like _ANNOTATABLE_TYPES above: _cuda_driver is None when cuda.bindings is # absent, so reading the enum at import time would break `import torch`. _NESTED_GRAPH_TYPES: set[Any] | None = None def _get_nested_graph_types() -> set[Any]: global _NESTED_GRAPH_TYPES if _NESTED_GRAPH_TYPES is None: node_types = _cuda_driver.CUgraphNodeType # pyrefly: ignore[missing-attribute] _NESTED_GRAPH_TYPES = { node_types.CU_GRAPH_NODE_TYPE_GRAPH, node_types.CU_GRAPH_NODE_TYPE_CONDITIONAL, } return _NESTED_GRAPH_TYPES # Pending scopes: (annotation, toolsIds discovered for the scope). # # Notably, this is NOT the current dynamic scope; instead, we fill this in # after we exit each mark_kernels region, and this holds the FULL extra # CUDA graph node attribution, steadily growing as we execute. # # Scopes are recorded as we exit a `mark_kernels`. Later, when we resolve # all of the annotations, we do merges on the annotations where # first-recorded-wins (it works: you end up preferring inner scopes over # outer scopes, what you'd expect). For example: # # with mark_kernels({"name": "outer", "color": "red"}): # with mark_kernels({"name": "inner"}): # y = x + 1 # toolsId: 101 # z = y * 2 # toolsId: 102 # # where the add captures kernel node with toolsId 101 and the mul 102. After # executing ALL of this code, the _pending_scopes would be (inner exits first, # so it is recorded first): # # ({"name": "inner"}, [101]) # ({"name": "outer", "color": "red"}, [101, 102]) # # toolsId 101 appears in both lists, and so the first-recorded-wins rule then # means we merge all the dicts to {"name": "inner", "color": "red"}. _pending_scopes: list[tuple[Any, list[int]]] = [] # TLS slot carrying the current annotation region: the annotations of every # open backward-annotating ``mark_kernels`` scope collapsed into one dict # (inner scopes win common keys). It lives in ThreadLocalPythonObjects rather # than a ContextVar or threading.local because that is the one Python-writable # TLS that ``at::ThreadLocalState`` snapshots into autograd engine worker and # device threads, so a backward bracket can re-establish the region for nodes # created while it executes (higher-order grad, checkpoint recomputation) even # across engine thread hops. Region dicts are never mutated after creation, so # holding a reference is a snapshot. _REGION_TLS_KEY = "cuda_graph_annotation_region" def _current_region() -> dict[str, Any] | None: if torch._C._is_key_in_tls(_REGION_TLS_KEY): return torch._C._get_obj_in_tls(_REGION_TLS_KEY) return None def _enter_region(collapsed: dict[str, Any]) -> dict[str, Any] | None: """Install ``collapsed`` as the current region; return the previous one.""" prev = _current_region() torch._C._stash_obj_in_tls(_REGION_TLS_KEY, collapsed) return prev def _exit_region(prev: dict[str, Any] | None) -> None: if prev is None: torch._C._remove_obj_from_tls(_REGION_TLS_KEY) else: torch._C._stash_obj_in_tls(_REGION_TLS_KEY, prev) # node.metadata marker recording that a node's backward hook pair is already # installed, so redundant firings of ``_freeze_region_hook`` (one per live # push, e.g. nested scopes) attach only once. An identity sentinel so no # other metadata user can collide with it. _HOOKED_KEY: Any = object() # Bumped by clear_kernel_annotations. Node hooks record only while the # generation they were created under is current, so clearing revokes # recording from every scope opened before the clear (the hooks stay on # the graph but become inert) without mutating live autograd graphs. _annotation_generation: int = 0 class _KernelScope(NamedTuple): """Capture-frontier snapshot taken at scope entry, consumed at scope exit.""" graph: Any frontier: list[Any] entry_root_keys: set[int] | None entry_direct_dependents: _ExistingDirectDependents def _begin_kernel_scope() -> _KernelScope | None: """Snapshot the current stream's capture frontier; ``None`` if not capturing.""" stream = _cuda_runtime.cudaStream_t( # pyrefly: ignore[missing-attribute] init_value=torch.cuda.current_stream().cuda_stream ) capture_state = _get_capture_state(stream) if capture_state is None: return None graph, frontier = capture_state entry_root_keys: set[int] | None = None entry_direct_dependents = { int(node): {int(dep) for dep in _get_dependent_nodes(node)} for node in frontier } if not frontier: entry_root_keys = {int(node) for node in _get_root_nodes(graph)} return _KernelScope(graph, frontier, entry_root_keys, entry_direct_dependents) def _end_kernel_scope(scope: _KernelScope) -> list[int]: """Walk nodes captured since ``scope`` was begun; return their toolsIds.""" if scope.frontier: scope_nodes = _collect_descendants( scope.frontier, existing_direct_dependents=scope.entry_direct_dependents, ) else: new_roots = [ node for node in _get_root_nodes(scope.graph) if int(node) not in (scope.entry_root_keys or set()) ] scope_nodes = _collect_descendants(new_roots, include_start_nodes=True) annotatable = _get_annotatable_types() nested = _get_nested_graph_types() nested_seen: set[str] = set() tools_ids: list[int] = [] for node in scope_nodes.values(): node_type = _get_node_type(node) if node_type in nested: nested_seen.add(node_type.name) if node_type not in annotatable: continue tools_ids.append( _check_cuda_bindings( _cuda_runtime.cudaGraphNodeGetToolsId( # pyrefly: ignore[missing-attribute] node ) ) ) if nested_seen: # The annotations recorded above are still correct -- the exec graph preserves # top-level node ids -- they just do not cover the nested body. warnings.warn( f"mark_kernels: this scope contains {sorted(nested_seen)} nodes; the work " "inside them is in a separate cudaGraph_t that this walk does not descend " "into, so it is left unannotated", stacklevel=4, ) return tools_ids class _BracketState(threading.local): """Per-thread stack of open backward brackets for one hook pair. Thread-local so concurrent executions of a retained node (separate graph tasks on different engine threads) cannot interleave each other's scope snapshots; within one thread pre/post pair up LIFO. If the node throws, its posthook never runs (the engine's call_function has no try/finally around fn) and the prehook's entry is left on the stack. That is harmless. Say a retained node runs three times on one thread and the second run throws: run 1: prehook push [e1]; posthook pops e1 -> [] run 2: prehook push [e2]; node throws, no pop -> [e2] run 3: prehook push [e2, e3]; posthook pops e3 -> [e2] Pops always take the top, so the leaked e2 sits below all later entries and every subsequent run still pops its own. The only loss is e2 itself: that run's kernels go unrecorded, which is moot since its backward failed. (The region/creation-hook TLS the prehook set is not leaked -- the engine task's ThreadLocalStateGuard restores those on unwind.) """ def __init__(self) -> None: # Entries are (scope, prev_region, merged_region) or _SKIPPED. self.stack: list[Any] = [] # Bracket-stack marker for "prehook did nothing": distinct from an entry with # scope=None, which means the region and creation hook were installed but no # capture was active (the posthook must still pop them). _SKIPPED: Any = object() def _freeze_region_hook(node: Any, generation: int) -> None: """Node creation hook: freeze the current region into ``node``'s bracket. Reads the region TLS slot and, if a region is open, installs the node's single backward hook pair closing over that snapshot. Every backward- annotating ``mark_kernels`` scope pushes one of these (and so does an executing hooked node, to cover nodes created during its backward), so a node created under nested scopes sees several identical firings; the ``node.metadata`` marker makes only the first attach. AccumulateGrad is always excluded: a leaf's node is created once and cached (possibly during warmup), so which scope would own it is an accident of first use, and its work (the ``.grad`` accumulation) belongs to the leaf rather than to any forward region. If accumulation ever needs annotation it should get a dedicated leaf-accumulation region, not first-use assignment. """ if isinstance(node, torch._C._functions.AccumulateGrad): return if _HOOKED_KEY in node.metadata: return frozen = _current_region() if frozen is None: return node.metadata[_HOOKED_KEY] = True _attach_backward_hooks(node, frozen, generation) def _attach_backward_hooks(node: Any, frozen: dict[str, Any], generation: int) -> None: """Register one pre/post hook pair bracketing ``node``'s backward execution. ``frozen`` is the region that was current when the node was created: the annotations of every scope then open, already collapsed inner-wins. The pair brackets the node's execution the same way ``mark_kernels`` brackets forward work: snapshot the capture frontier before, collect the new nodes after. It runs on the autograd engine thread executing the node, which during a whole-graph capture participates in the capture; when backward runs outside a capture the prehook sees no active capture and the pair is a no-op. Executing the node is semantically re-entering its frozen region, so the prehook merges ``frozen`` over the live region (a ``mark_kernels`` enclosing the whole ``backward()`` call, if any) and installs the result while the node runs. That gives nodes created during its backward (``create_graph=True``, checkpoint recomputation) the forward region's ownership -- a creation hook is pushed alongside so they get hooked even when this backward runs eagerly and no scope's hook is live -- and it puts a ``mark_kernels`` opened inside the execution dynamically inner, so it outranks ``frozen``, which outranks the enclosing scope. The posthook records the merged region once; scope completion order then yields exactly that precedence at resolve. The pops are exception-safe without a finally: each engine task runs under an ``at::ThreadLocalStateGuard`` that restores both the creation-hook TLS and the region slot even when the node throws. ``generation`` is the annotation generation the owning scope was opened under; the bracket only records (and only propagates to created nodes) while it is current, so ``clear_kernel_annotations`` makes stale hooks inert. """ state = _BracketState() def creation_hook(child: Any) -> None: _freeze_region_hook(child, generation) def prehook(_grad_outputs: Any) -> None: if generation != _annotation_generation or _is_tools_id_unavailable(): state.stack.append(_SKIPPED) return # Re-establish ownership even when this backward is not itself # captured: kernels of nodes created here may be captured by a # later (e.g. second-order) annotated capture. merged = {**(_current_region() or {}), **frozen} prev = _enter_region(merged) torch._C._autograd._push_node_creation_hook(creation_hook) scope = _begin_kernel_scope() if _annotations_enabled else None state.stack.append((scope, prev, merged)) def posthook(_grad_inputs: Any, _grad_outputs: Any) -> None: if not state.stack: return entry = state.stack.pop() if entry is _SKIPPED: return scope, prev, merged = entry torch._C._autograd._pop_node_creation_hook() _exit_region(prev) if scope is None: return tools_ids = _end_kernel_scope(scope) if tools_ids: tagged = {**merged, "autograd_phase": "backward"} _pending_scopes.append((tagged, tools_ids)) node.register_prehook(prehook) node.register_hook(posthook) @contextmanager # type: ignore[arg-type] def mark_kernels(annotation: str | dict[str, Any], *, backward: bool = True): r"""mark_kernels(annotation, *, backward=True) Context manager that annotates GPU work captured within its scope. Must be used inside an active :class:`torch.cuda.graph` capture with ``enable_annotations=True``. Every kernel, memcpy, and memset node the capture adds within the scope is tagged with :attr:`annotation`. Outside a capture, with annotations disabled, or when :func:`is_available` is ``False``, the context manager is a no-op. When scopes overlap on the same node (e.g. nested scopes), their annotation dicts are merged key-by-key with the inner scope winning common keys. By default backward work is annotated too: autograd nodes created by forward operations inside the scope get hooks (via :class:`torch.autograd.graph.node_creation_hook`) that bracket their backward execution, so when the backward pass is itself captured -- in the same capture as the forward or in a later one -- its kernels are tagged with the same annotation, plus an ``"autograd_phase": "backward"`` key marking them as backward work (``"autograd_phase"`` is therefore reserved: backward annotation overwrites it). When backward runs outside a capture the hooks are a no-op. Ownership extends to higher-order gradients: nodes created while a hooked node executes (``create_graph=True``, checkpoint recomputation) inherit its annotations, so a later grad-of-grad capture is attributed too. ``AccumulateGrad`` nodes are never annotated: a leaf's node is created once and cached, so scope ownership would be an accident of first use. Pass ``backward=False`` to annotate only the forward work, e.g. when a wrapper implements its own backward attribution. The keyword's presence also serves as the feature probe for that native backward support: ``"backward" in inspect.signature(mark_kernels).parameters``. Implementation: on entry, records the current stream's capture frontier and its existing direct dependents; on scope exit, walks only the dependent nodes added since entry (falling back to newly created graph roots when the scope is the first captured work). Args: annotation (str or dict): Metadata to attach to each captured node. A string ``s`` is recorded as ``{"name": s}``. Dict values must be picklable. The key ``"name"`` names the region in trace tooling; ``"stream"`` is reserved for stream-lane assignment. backward (bool): Whether to also annotate the backward kernels of autograd nodes created inside the scope. Default: ``True``. .. note:: The nodes to annotate must be reachable from the capture frontier of the stream that is current on scope entry. Work on a different already-capturing stream must be synchronized with the current stream first. .. note:: Child-graph and conditional nodes have bodies in a separate ``cudaGraph_t`` that this walk does not descend into, so their work is left unannotated and a warning is issued. Descending is possible (``cudaGraphNodeGetParams`` exposes the body graphs), but would not be enough on its own: a body's nodes are numbered in that graph's id space and are renumbered again when the exec graph inlines them, and nothing exposes that renumbering, so :func:`remap_to_exec_graph` could not key the annotations to what a profiler reports. For the same reason a scope *inside* a conditional body (``torch.cond`` / ``torch.while_loop``) records nothing at all. .. warning:: This API is in prototype and may change in future releases. Example:: >>> # xdoctest: +SKIP("requires cuda-bindings and driver >= 13.1") >>> g = torch.cuda.CUDAGraph() >>> x = torch.randn(8, device="cuda") >>> with torch.cuda.graph(g, enable_annotations=True): ... with torch.cuda.graph_annotations.mark_kernels("phase_A"): ... y = x + 1 """ if not _annotations_enabled or _is_tools_id_unavailable(): yield return if isinstance(annotation, str): annotation = {"name": annotation} scope = _begin_kernel_scope() if scope is None: yield return if ( _capture_root_graph_id is not None and _graph_id(scope.graph) != _capture_root_graph_id ): # Inside a conditional node's body: torch.cond / torch.while_loop capture into a # separate cudaGraph_t. Its node ids are in that graph's id space and are # renumbered again in the exec graph, so anything recorded here would be a key # that matches nothing in a trace. Record nothing rather than dead keys. warnings.warn( "mark_kernels: this scope is inside a CUDA graph conditional-node body " "(torch.cond / torch.while_loop), which is captured into a separate " "cudaGraph_t whose node ids are never remapped to the exec graph; " "nothing is annotated for it", stacklevel=3, ) yield return generation = _annotation_generation def creation_hook(node: Any) -> None: _freeze_region_hook(node, generation) # Enter the region even when backward=False: the region must reflect the # full dynamic scope so that nodes hooked by a nested backward=True scope # (or by an executing hooked node) freeze this annotation too. prev = _enter_region({**(_current_region() or {}), **annotation}) try: if backward: with torch.autograd.graph.node_creation_hook(creation_hook): yield else: yield finally: _exit_region(prev) tools_ids = _end_kernel_scope(scope) if tools_ids: _pending_scopes.append((annotation, tools_ids)) def maybe_stamp_capture_graph_id(torch_cuda_graph: torch.cuda.CUDAGraph) -> None: """Record the captured graph's id on the graph object for later remap. Called from ``capture_end`` in the window after capture ends and before the graph is finalized, where the template ``cudaGraph_t`` is live for both ``keep_graph`` modes (and reachable via ``raw_cuda_graph``). ``remap_to_exec_graph`` later matches this graph's annotations to its exec id via this stamp, without relying on call ordering. No-op if annotations are disabled or ``cudaGraphNodeGetToolsId`` is unavailable. Harmless when no ``mark_kernels`` regions run: the annotation map stays empty so resolve and remap are no-ops. """ if not _annotations_enabled or _is_tools_id_unavailable(): return # Past the _is_tools_id_unavailable() guard cuda-bindings is present and the # driver supports the toolsId API (same version gate as cudaGraphGetId), so # any error here is unexpected: error-check and let it raise. cudaGraphGetId # accepts the raw cudaGraph_t handle (int) directly. torch_cuda_graph._capture_graph_id = _check_cuda_bindings( _cuda_runtime.cudaGraphGetId( # pyrefly: ignore[missing-attribute] torch_cuda_graph.raw_cuda_graph() ) ) # Fresh capture: annotations are keyed by this capture id until remapped. torch_cuda_graph._remapped_exec_id = None def resolve_pending_annotations() -> None: """Resolve pending scope toolsIds into kernel annotations.""" if not _pending_scopes: return try: per_tools_id: defaultdict[int, list[Any]] = defaultdict(list) for annotation, tools_ids in _pending_scopes: for tools_id in tools_ids: per_tools_id[tools_id].append(annotation) for tools_id, ann_list in per_tools_id.items(): if len(ann_list) == 1: _kernel_annotations[tools_id].append(ann_list[0]) continue merged: dict[str, Any] = {} for annotation in ann_list: if isinstance(annotation, dict): for key, value in annotation.items(): merged.setdefault(key, value) else: merged.setdefault("name", annotation) _kernel_annotations[tools_id].append(merged) except Exception: logger.exception("resolve_pending_annotations failed") finally: _pending_scopes.clear() def remap_to_exec_graph(torch_cuda_graph: torch.cuda.CUDAGraph) -> None: """Remap annotation keys from capture graph ID to exec graph ID. During capture, toolsId encodes the capture graph's ID in the upper 32 bits. After instantiation, the profiler uses the exec graph's ID. This function rewrites the keys so annotations match the trace. The graph's capture id is read from the ``_capture_graph_id`` stamped on it by ``maybe_stamp_capture_graph_id`` in the capture_end window, so only the annotations belonging to this graph are rekeyed. This is order-independent and correct when several graphs are captured in sequence: call once per graph. Graphs captured with annotations disabled have no capture id and are skipped. The exec graph id is only defined once the graph is instantiated. With ``keep_graph=True`` instantiation is deferred past the ``torch.cuda.graph()`` context, so the remap is driven from the graph's ``instantiate()``/ ``replay()`` instead of context exit. Each ``instantiate()`` (even on an unmodified template) produces a fresh exec graph id, so this rekeys from wherever the annotations are currently keyed -- the capture id before the first remap, the previous exec id after a re-instantiate -- to the current exec id. ``_remapped_exec_id`` tracks that current key; when it already matches the live exec id (e.g. replay after instantiate) this is a no-op. """ capture_graph_id = torch_cuda_graph._capture_graph_id if not _kernel_annotations or capture_graph_id is None: return exec_graph_id = _check_cuda_bindings( _cuda_runtime.cudaGraphExecGetId( # pyrefly: ignore[missing-attribute] torch_cuda_graph.raw_cuda_graph_exec() ) ) current_key_id = ( capture_graph_id if torch_cuda_graph._remapped_exec_id is None else torch_cuda_graph._remapped_exec_id ) if current_key_id == exec_graph_id: return remapped = _rekey_annotations(_kernel_annotations, current_key_id, exec_graph_id) _kernel_annotations.clear() _kernel_annotations.update(remapped) torch_cuda_graph._remapped_exec_id = exec_graph_id def _rekey_annotations( annotations: dict[int, list[Any]], capture_graph_id: int, exec_graph_id: int, ) -> dict[int, list[Any]]: """Rekey one graph's annotations from its capture id to its exec id. A toolsId packs the graph id in the upper 32 bits and the node id in the lower 32. Only entries whose upper bits match ``capture_graph_id`` are rewritten to ``exec_graph_id``; entries from other graphs (already remapped to their own exec ids, or pending their own remap) are kept as-is. When two capture-side ids collide on the same rekeyed id their lists are merged. """ remapped: dict[int, list[Any]] = {} for tools_id, ann_list in annotations.items(): if tools_id >> 32 != capture_graph_id: remapped[tools_id] = ann_list continue node_id = tools_id & 0xFFFFFFFF new_tools_id = (exec_graph_id << 32) | node_id if new_tools_id in remapped: remapped[new_tools_id].extend(ann_list) else: remapped[new_tools_id] = list(ann_list) return remapped def resolve_and_remap(torch_cuda_graph: torch.cuda.CUDAGraph) -> None: """Resolve any pending scopes and remap one graph in a single call. Shorthand for ``resolve_pending_annotations()`` followed by ``remap_to_exec_graph(graph)``; the pair normally run after a capture. """ resolve_pending_annotations() remap_to_exec_graph(torch_cuda_graph) def get_kernel_annotations() -> Mapping[int, list[Any]]: r"""get_kernel_annotations() -> Mapping[int, list] Return the live registry of recorded kernel annotations. Keys are opaque integers matching the ``graph node id`` field that CUPTI-based profilers attach to kernel events; values are the lists of annotation dicts recorded for that node. The registry accumulates across captures and is global to the process. The returned mapping is a **live view**: it is updated in place when a graph is instantiated (annotation keys are rekeyed to the executable graph's ids), so a reference obtained early stays current. Keys are valid for joining against a profiler trace once the corresponding graphs have been instantiated. Treat the mapping as read-only; snapshot it with ``dict(...)`` if isolation is needed. .. warning:: This API is in prototype and may change in future releases. Example:: >>> # xdoctest: +SKIP("requires cuda-bindings and driver >= 13.1") >>> annotations = torch.cuda.graph_annotations.get_kernel_annotations() >>> with open("annotations.pkl", "wb") as f: ... pickle.dump(dict(annotations), f) """ return _kernel_annotations def clear_kernel_annotations() -> None: r"""clear_kernel_annotations() -> None Clear all recorded kernel annotations. The annotation registry is process-global and accumulates across captures; long-running workloads that capture many graphs should clear it once recorded annotations have been consumed (e.g. after saving them alongside a profiler trace). Clearing forgets everything recorded so far and revokes recording from every scope opened before the clear: backward hooks that :func:`mark_kernels` attached to existing autograd nodes stay on the graph but become inert. Scopes opened after the clear record normally. .. warning:: This API is in prototype and may change in future releases. """ global _annotation_generation _kernel_annotations.clear() _pending_scopes.clear() _annotation_generation += 1 def remove_kernel_annotations(exec_graph_ids: Iterable[int]) -> None: """Drop kernel-annotation entries whose exec graph id (tools_id >> 32) is in exec_graph_ids, so the map does not grow across the run. Run by the annotation resolver's graph-destroy handler.""" ids = set(exec_graph_ids) if not ids: return for key in [k for k in _kernel_annotations if k >> 32 in ids]: del _kernel_annotations[key] # Counter-based stream ID registry. IDs start at 60 (above the highest # observed non-graphed CUDA stream ID) so every assigned lane is visually # distinct in Perfetto and doesn't collide with real streams. _stream_id_counter: int = 60 _stream_id_map: dict[int, int] = {} def _get_stream_id(stream: torch.cuda.Stream) -> int: """Return a small, stable stream ID for the given CUDA stream.""" global _stream_id_counter key = stream.cuda_stream if key not in _stream_id_map: _stream_id_map[key] = _stream_id_counter _stream_id_counter += 1 return _stream_id_map[key] def get_stream_for_pg(pg_key: str) -> int: """Return a unique stream ID for the given process group key.""" global _stream_id_counter if pg_key not in _stream_id_map: _stream_id_map[pg_key] = _stream_id_counter # type: ignore[assignment] _stream_id_counter += 1 return _stream_id_map[pg_key] # type: ignore[return-value] @contextmanager # type: ignore[arg-type] def mark_stream(stream: torch.cuda.Stream, annotation: str | dict[str, Any]): """Switch to stream, inject its ID into annotation, and mark kernels. If *stream* is already the current stream, no stream switch or stream ID injection happens — the kernels stay on whatever stream is active (which keeps the trace faithful when e.g. FSDP uses the current stream for copy-in instead of a separate one). When switching to a different stream, this snapshots the current capturing branch before the target stream runs marked work. If the target stream is already capturing, the marked work must still be synchronized with the current stream so it is reachable from that snapped frontier. """ if not _annotations_enabled: with torch.cuda.stream(stream): yield return if stream.cuda_stream == torch.cuda.current_stream().cuda_stream: with mark_kernels(annotation): yield else: if isinstance(annotation, str): annotation = {"name": annotation} if isinstance(annotation, dict): annotation["stream"] = _get_stream_id(stream) with mark_kernels(annotation): with torch.cuda.stream(stream): yield