|
| 1 | +# Copyright (c) 2025, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED. |
| 2 | +# |
| 3 | +# SPDX-License-Identifier: LicenseRef-NVIDIA-SOFTWARE-LICENSE |
| 4 | + |
| 5 | +from __future__ import annotations |
| 6 | + |
| 7 | +import os |
| 8 | +import warnings |
| 9 | +import weakref |
| 10 | +from dataclasses import dataclass |
| 11 | +from typing import TYPE_CHECKING, Optional, Tuple, Union |
| 12 | + |
| 13 | +if TYPE_CHECKING: |
| 14 | + from cuda.core.experimental._stream import Stream |
| 15 | +from cuda.core.experimental._utils import driver, handle_return, precondition |
| 16 | + |
| 17 | + |
| 18 | +@dataclass |
| 19 | +class DebugPrintOptions: |
| 20 | + """ |
| 21 | + """ |
| 22 | + |
| 23 | + VERBOSE: bool = False |
| 24 | + RUNTIME_TYPES: bool = False |
| 25 | + KERNEL_NODE_PARAMS: bool = False |
| 26 | + MEMCPY_NODE_PARAMS: bool = False |
| 27 | + MEMSET_NODE_PARAMS: bool = False |
| 28 | + HOST_NODE_PARAMS: bool = False |
| 29 | + EVENT_NODE_PARAMS: bool = False |
| 30 | + EXT_SEMAS_SIGNAL_NODE_PARAMS: bool = False |
| 31 | + EXT_SEMAS_WAIT_NODE_PARAMS: bool = False |
| 32 | + KERNEL_NODE_ATTRIBUTES: bool = False |
| 33 | + HANDLES: bool = False |
| 34 | + MEM_ALLOC_NODE_PARAMS: bool = False |
| 35 | + MEM_FREE_NODE_PARAMS: bool = False |
| 36 | + BATCH_MEM_OP_NODE_PARAMS: bool = False |
| 37 | + EXTRA_TOPO_INFO: bool = False |
| 38 | + CONDITIONAL_NODE_PARAMS: bool = False |
| 39 | + |
| 40 | + |
| 41 | +class GraphBuilder: |
| 42 | + """TBD |
| 43 | +
|
| 44 | + Directly creating a :obj:`~_graph.GraphBuilder` is not supported due |
| 45 | + to ambiguity. New graph builders should instead be created through a |
| 46 | + :obj:`~_device.Device`, or a :obj:`~_stream.stream` object |
| 47 | + """ |
| 48 | + |
| 49 | + def __init__(self): |
| 50 | + raise NotImplementedError( |
| 51 | + "directly creating a Graph object can be ambiguous. Please either " |
| 52 | + "call Device.create_graph() or stream.creating_graph()" |
| 53 | + ) |
| 54 | + |
| 55 | + @staticmethod |
| 56 | + def _init(stream, is_root=True): |
| 57 | + self = GraphBuilder.__new__(GraphBuilder) |
| 58 | + # TODO: I need to know if we own this stream object. |
| 59 | + # If from Device(), then we can destroy it on close |
| 60 | + # If from Stream, then we can't |
| 61 | + self._stream = stream |
| 62 | + self._capturing = False |
| 63 | + self._is_root = is_root # TODO: Is this info needed? |
| 64 | + return self |
| 65 | + |
| 66 | + def _check_capture_stream_provided(self, *args, **kwargs): |
| 67 | + if self._stream == None: |
| 68 | + raise RuntimeError("Tried to use a stream capture operation on a graph builder without a stream") |
| 69 | + |
| 70 | + @property |
| 71 | + def legacy_stream_capture(self) -> Stream: |
| 72 | + return self._stream |
| 73 | + |
| 74 | + @property |
| 75 | + def is_root_builder(self) -> bool: |
| 76 | + return self._is_root_builder |
| 77 | + |
| 78 | + @precondition(_check_capture_stream_provided) |
| 79 | + def begin_capture(self, mode="global"): |
| 80 | + # Supports "global", "local" or "relaxed" |
| 81 | + # TODO; Test case for each mode and fail |
| 82 | + if mode == "global": |
| 83 | + capture_mode = driver.CUstreamCaptureMode.CU_STREAM_CAPTURE_MODE_GLOBAL |
| 84 | + elif mode == "local": |
| 85 | + capture_mode = driver.CUstreamCaptureMode.CU_STREAM_CAPTURE_MODE_THREAD_LOCAL |
| 86 | + elif mode == "relaxed": |
| 87 | + capture_mode = driver.CUstreamCaptureMode.CU_STREAM_CAPTURE_MODE_RELAXED |
| 88 | + else: |
| 89 | + raise ValueError(f"Only 'global', 'local' or 'relaxed' capture mode are supported, got {capture_mode}") |
| 90 | + |
| 91 | + handle_return(driver.cuStreamBeginCapture(self._stream.handle, capture_mode)) |
| 92 | + self._capturing = True |
| 93 | + |
| 94 | + @precondition(_check_capture_stream_provided) |
| 95 | + def is_capture_active(self) -> bool: |
| 96 | + result = handle_return(driver.cuStreamGetCaptureInfo(self._stream.handle)) |
| 97 | + |
| 98 | + capture_status = result[0] |
| 99 | + if capture_status == driver.CUstreamCaptureStatus.CU_STREAM_CAPTURE_STATUS_NONE: |
| 100 | + return False |
| 101 | + elif capture_status == driver.CUstreamCaptureStatus.CU_STREAM_CAPTURE_STATUS_ACTIVE: |
| 102 | + return True |
| 103 | + elif capture_status == driver.CUstreamCaptureStatus.CU_STREAM_CAPTURE_STATUS_INVALIDATED: |
| 104 | + raise RuntimeError("Stream is part of a capture sequence that has been invalidated, but " |
| 105 | + "not terminated. The capture sequence must be terminated with self.`().") |
| 106 | + else: |
| 107 | + raise NotImplementedError(f"Unsupported capture stuse type received: {capture_status}") |
| 108 | + |
| 109 | + @precondition(_check_capture_stream_provided) |
| 110 | + def end_capture(self): |
| 111 | + if not self._capturing: |
| 112 | + raise RuntimeError("Stream is not capturing") |
| 113 | + |
| 114 | + self._graph = handle_return(driver.cuStreamEndCapture(self._stream.handle)) |
| 115 | + self._capturing = False |
| 116 | + |
| 117 | + def debug_dot_print(self, path, options: Optional[DebugPrintOptions] = None): |
| 118 | + # TODO: We should be able to print one while the capture is happening right? Just need to make sure driver version is new enough. |
| 119 | + if self._graph == None: |
| 120 | + raise RuntimeError("Graph needs to be built before generating a DOT debug file") |
| 121 | + |
| 122 | + # TODO: Apply each option to the value |
| 123 | + options_value = 0 |
| 124 | + |
| 125 | + handle_return(driver.cuGraphDebugDotPrint(self._graph, path, options_value)) |
| 126 | + |
| 127 | + def split(self, count) -> Tuple[Graph, ...]: |
| 128 | + if count <= 1: |
| 129 | + raise ValueError(f"Invalid split count: expecting >= 2, got {count}") |
| 130 | + |
| 131 | + # 1. Record an event on our stream |
| 132 | + event = self._stream.record() |
| 133 | + |
| 134 | + # TODO: Steps 2,3,4 can be combined under a single loop |
| 135 | + |
| 136 | + # 2. Create a streams for each of the new splits |
| 137 | + # TODO: Optimization where one of the split stream is allowed to use |
| 138 | + # TODO: Should use the same stream options as initial stream?? |
| 139 | + split_stream = [self._stream.device.create_stream() for i in range(count)] |
| 140 | + |
| 141 | + # 3. Have each new stream wait on our singular event |
| 142 | + for stream in split_stream: |
| 143 | + stream.wait(event) |
| 144 | + |
| 145 | + # 4. Discard the event |
| 146 | + # TODO: Is this actually allowed when using with a graph? Surely, since it just needs to create an edge for us... right? |
| 147 | + event.close() |
| 148 | + |
| 149 | + # 5. Create new graph builders for each new stream split |
| 150 | + return [GraphBuilder._init(stream=stream, is_root=False) for stream in split_stream] |
| 151 | + |
| 152 | + def join(self, *graph_builders): |
| 153 | + if len(graph_builders) < 1: |
| 154 | + raise ValueError("Must specify which graphs should join but none were given") |
| 155 | + |
| 156 | + for graph in graph_builders: |
| 157 | + self._stream.wait(graph.legacy_stream_capture) |
| 158 | + # TODO: Can we close each of those new streams? Do we need weakref? |
| 159 | + graph.close() |
| 160 | + |
| 161 | + def close(self): |
| 162 | + if self._capturing: |
| 163 | + raise RuntimeError("Trying to close a graph builder who is still capturing") |
| 164 | + # Can we call this directly? Or relying on weakref enough? |
| 165 | + self._stream.close() |
| 166 | + |
| 167 | + |
| 168 | +class Graph: |
| 169 | + """ |
| 170 | + """ |
| 171 | + |
| 172 | + def __init__(self): |
| 173 | + raise RuntimeError("directly constructing a Graph instance is not supported") |
| 174 | + |
| 175 | + @staticmethod |
| 176 | + def _init(graph): |
| 177 | + self = Graph.__new__(Graph) |
| 178 | + self._graph = graph |
| 179 | + return self |
0 commit comments