Coverage for cuda/core/graph/_graph_builder.pyx: 78.08%
520 statements
« prev ^ index » next coverage.py v7.17.2, created at 2026-11-05 01:11 +0000
1# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
2#
3# SPDX-License-Identifier: Apache-2.0
5from dataclasses import dataclass
6from typing import TYPE_CHECKING
8from libc.stdint cimport intptr_t
10from cuda.bindings cimport cydriver
12from cuda.core.graph._graph_definition cimport (
13 GraphCondition,
14 GraphDefinition,
15 GD_check_valid,
16)
17from cuda.core.graph._graph_node cimport GraphNode, GN_check_valid
18from cuda.core._host_callback cimport _resolve_host_callback
19from cuda.core.graph._subclasses cimport (
20 ChildGraphNode,
21 EventRecordNode,
22 EventWaitNode,
23 ExecutableChildGraphNode,
24 ExecutableEventRecordNode,
25 ExecutableEventWaitNode,
26 ExecutableGraphNode,
27 ExecutableHostCallbackNode,
28 ExecutableKernelNode,
29 ExecutableMemcpyNode,
30 ExecutableMemsetNode,
31 HostCallbackNode,
32 KernelNode,
33 MemcpyNode,
34 MemsetNode,
35 create_executable_node_view,
36)
37from cuda.core._rt cimport attach_rollback_failure, report_cuda_error
38from cuda.core._rt cimport (
39 GraphExecHandle,
40 GraphHandle,
41 OpaqueHandle,
42 PreparedAttachment,
43 as_cu, as_py,
44 create_child_graph_handle, create_graph_exec_handle, create_graph_handle,
45 get_last_error,
46 graph_clone_attachments,
47 graph_commit_attachment,
48 graph_exec_update,
49 graph_prepare_attachment,
50 invalidate_child_graph_state,
51 invalidate_root_graph_state,
52 retry_deferred_cleanup,
53)
54from cuda.core._stream cimport Stream, Stream_accept
55from cuda.core._utils.cuda_utils cimport HANDLE_RETURN
56from cuda.core._utils.version cimport cy_driver_version
58from cuda.core._utils.cuda_utils import (
59 CUDAError,
60 driver,
61 handle_return,
62)
64if TYPE_CHECKING:
65 from collections.abc import Callable
66 from typing import overload
68 from cuda.core._utils.pycompat import BufferProtocol
70 from cuda.core.graph._graph_definition import GraphDefinition
71else:
72 # Cython applies a method decorator by rebinding the type's dict entry, so
73 # typing.overload would replace the compiled ``__getitem__`` slot wrapper
74 # with its placeholder. The overloads only serve the stub; at run time the
75 # decorator must leave the method alone.
76 def overload(f):
77 return f
79# Each ``@overload`` stub of ``Graph.__getitem__`` still compiles to a static
80# wrapper function that nothing references, because only the final definition
81# is bound. gcc and clang report those under -Wunused-function, so this file
82# turns that warning off for its own translation unit. The guard keeps the
83# pragma away from MSVC, whose equivalent warning is off at the default level.
84cdef extern from *:
85 """
86 #if defined(__GNUC__)
87 #pragma GCC diagnostic ignored "-Wunused-function"
88 #endif
89 """
91__all__ = ['Graph', 'GraphBuilder', 'GraphCompleteOptions', 'GraphDebugPrintOptions']
94@dataclass
95class GraphDebugPrintOptions:
96 """Options for debug_dot_print().
98 Attributes
99 ----------
100 verbose : bool
101 Output all debug data as if every debug flag is enabled (Default to False)
102 runtime_types : bool
103 Use CUDA Runtime structures for output (Default to False)
104 kernel_node_params : bool
105 Adds kernel parameter values to output (Default to False)
106 memcpy_node_params : bool
107 Adds memcpy parameter values to output (Default to False)
108 memset_node_params : bool
109 Adds memset parameter values to output (Default to False)
110 host_node_params : bool
111 Adds host parameter values to output (Default to False)
112 event_node_params : bool
113 Adds event parameter values to output (Default to False)
114 ext_semas_signal_node_params : bool
115 Adds external semaphore signal parameter values to output (Default to False)
116 ext_semas_wait_node_params : bool
117 Adds external semaphore wait parameter values to output (Default to False)
118 kernel_node_attributes : bool
119 Adds kernel node attributes to output (Default to False)
120 handles : bool
121 Adds node handles and every kernel function handle to output (Default to False)
122 mem_alloc_node_params : bool
123 Adds memory alloc parameter values to output (Default to False)
124 mem_free_node_params : bool
125 Adds memory free parameter values to output (Default to False)
126 batch_mem_op_node_params : bool
127 Adds batch mem op parameter values to output (Default to False)
128 extra_topo_info : bool
129 Adds edge numbering information (Default to False)
130 conditional_node_params : bool
131 Adds conditional node parameter values to output (Default to False)
133 """
135 verbose: bool = False
136 runtime_types: bool = False
137 kernel_node_params: bool = False
138 memcpy_node_params: bool = False
139 memset_node_params: bool = False
140 host_node_params: bool = False
141 event_node_params: bool = False
142 ext_semas_signal_node_params: bool = False
143 ext_semas_wait_node_params: bool = False
144 kernel_node_attributes: bool = False
145 handles: bool = False
146 mem_alloc_node_params: bool = False
147 mem_free_node_params: bool = False
148 batch_mem_op_node_params: bool = False
149 extra_topo_info: bool = False
150 conditional_node_params: bool = False
152 def _to_flags(self) -> int:
153 """Convert options to CUDA driver API flags (internal use)."""
154 flags = 0 2jda
155 if self.verbose: 2jda
156 flags |= driver.CUgraphDebugDot_flags.CU_GRAPH_DEBUG_DOT_FLAGS_VERBOSE 2jda
157 if self.runtime_types: 2jda
158 flags |= driver.CUgraphDebugDot_flags.CU_GRAPH_DEBUG_DOT_FLAGS_RUNTIME_TYPES 1a
159 if self.kernel_node_params: 2jda
160 flags |= driver.CUgraphDebugDot_flags.CU_GRAPH_DEBUG_DOT_FLAGS_KERNEL_NODE_PARAMS 1a
161 if self.memcpy_node_params: 2jda
162 flags |= driver.CUgraphDebugDot_flags.CU_GRAPH_DEBUG_DOT_FLAGS_MEMCPY_NODE_PARAMS 1a
163 if self.memset_node_params: 2jda
164 flags |= driver.CUgraphDebugDot_flags.CU_GRAPH_DEBUG_DOT_FLAGS_MEMSET_NODE_PARAMS 1a
165 if self.host_node_params: 2jda
166 flags |= driver.CUgraphDebugDot_flags.CU_GRAPH_DEBUG_DOT_FLAGS_HOST_NODE_PARAMS 1a
167 if self.event_node_params: 2jda
168 flags |= driver.CUgraphDebugDot_flags.CU_GRAPH_DEBUG_DOT_FLAGS_EVENT_NODE_PARAMS 1a
169 if self.ext_semas_signal_node_params: 2jda
170 flags |= driver.CUgraphDebugDot_flags.CU_GRAPH_DEBUG_DOT_FLAGS_EXT_SEMAS_SIGNAL_NODE_PARAMS 1a
171 if self.ext_semas_wait_node_params: 2jda
172 flags |= driver.CUgraphDebugDot_flags.CU_GRAPH_DEBUG_DOT_FLAGS_EXT_SEMAS_WAIT_NODE_PARAMS 1a
173 if self.kernel_node_attributes: 2jda
174 flags |= driver.CUgraphDebugDot_flags.CU_GRAPH_DEBUG_DOT_FLAGS_KERNEL_NODE_ATTRIBUTES 1a
175 if self.handles: 2jda
176 flags |= driver.CUgraphDebugDot_flags.CU_GRAPH_DEBUG_DOT_FLAGS_HANDLES 2jda
177 if self.mem_alloc_node_params: 2jda
178 flags |= driver.CUgraphDebugDot_flags.CU_GRAPH_DEBUG_DOT_FLAGS_MEM_ALLOC_NODE_PARAMS 1a
179 if self.mem_free_node_params: 2jda
180 flags |= driver.CUgraphDebugDot_flags.CU_GRAPH_DEBUG_DOT_FLAGS_MEM_FREE_NODE_PARAMS 1a
181 if self.batch_mem_op_node_params: 2jda
182 flags |= driver.CUgraphDebugDot_flags.CU_GRAPH_DEBUG_DOT_FLAGS_BATCH_MEM_OP_NODE_PARAMS 1a
183 if self.extra_topo_info: 2jda
184 flags |= driver.CUgraphDebugDot_flags.CU_GRAPH_DEBUG_DOT_FLAGS_EXTRA_TOPO_INFO 1a
185 if self.conditional_node_params: 2jda
186 flags |= driver.CUgraphDebugDot_flags.CU_GRAPH_DEBUG_DOT_FLAGS_CONDITIONAL_NODE_PARAMS 1a
187 return flags 2jda
190@dataclass
191class GraphCompleteOptions:
192 """Options for graph instantiation.
194 Attributes
195 ----------
196 auto_free_on_launch : bool, optional
197 Automatically free memory allocated in a graph before relaunching. (Default to False)
198 upload_stream : Stream, optional
199 Stream to use to automatically upload the graph after completion. (Default to None)
200 device_launch : bool, optional
201 Configure the graph to be launchable from the device. This flag can only
202 be used on platforms which support unified addressing. This flag cannot be
203 used in conjunction with auto_free_on_launch. (Default to False)
204 use_node_priority : bool, optional
205 Run the graph using the per-node priority attributes rather than the
206 priority of the stream it is launched into. (Default to False)
208 """
210 auto_free_on_launch: bool = False
211 upload_stream: Stream | None = None
212 device_launch: bool = False
213 use_node_priority: bool = False
216def _instantiate_graph(source, options: GraphCompleteOptions | None = None) -> Graph:
217 cdef GraphHandle h_graph
218 cdef GraphExecHandle h_exec
219 cdef cydriver.CUresult status
221 if isinstance(source, GraphBuilder): 2P nbF L M N O cbH : x z 3 y odG ; 5 = 6 4 b c d e f g h i j k l m n o p q A r s t u v B C D E +b,b-b.b/bSb:bTb#b;b=b?b@bUb[bVb$bWbXb]bYb%bZb^b0b'bwcbdcdxcddedfd5bscmdgdyczcAc_b`bhdBctcucCcDcEcFcGcHcvcIcJcKcLcMc{bNcOcPcQcRcScTc| ^ ` gbabdb} _ { hbbbeb+ 7 $ , 8 % 0 V W - 9 ' . ! ( ac|bUcVcWcXcYcZc0c1c2c3c4c5c6cbc:c;c=c?c@c2 [c]c^c/cccdcecfcgchcicjckclcmcncoc_cpc7c8c9c!c#c$c%c'c(c)c*c+c,cRb2b-c.cR Q X Y / qcrc3b6b7bw 1 (b)b) ? ] }b~bfb* @ K I ~ [
222 GB_check_open(<GraphBuilder>source) 2P F L M N O cbH : x z 3 y G ; 5 = 6 4 b c d e f g h i j k l m n o p q A r s t u v B C D E | ^ ` gbabdb} _ { hbbbeb+ 7 $ , 8 % 0 V W - 9 ' . ! ( w 1 ) ? ] fb* @ K I ~ [
223 h_graph = (<GraphBuilder>source)._h_graph 2P F L M N O cbH : x z 3 y G ; 5 = 6 4 b c d e f g h i j k l m n o p q A r s t u v B C D E | ^ ` gbabdb} _ { hbbbeb+ 7 $ , 8 % 0 V W - 9 ' . ! ( w 1 ) ? ] fb* @ K I ~ [
224 elif isinstance(source, GraphDefinition): 2nbod+b,b-b.b/bSb:bTb#b;b=b?b@bUb[bVb$bWbXb]bYb%bZb^b0b'bwcbdcdxcddedfd5bscmdgdyczcAc_b`bhdBctcucCcDcEcFcGcHcvcIcJcKcLcMc{bNcOcPcQcRcScTcac|bUcVcWcXcYcZc0c1c2c3c4c5c6cbc:c;c=c?c@c2 [c]c^c/cccdcecfcgchcicjckclcmcncoc_cpc7c8c9c!c#c$c%c'c(c)c*c+c,cRb2b-c.cR Q X Y / qcrc3b6b7b(b)b}b~b
225 GD_check_valid(<GraphDefinition>source) 2nbod+b,b-b.b/bSb:bTb#b;b=b?b@bUb[bVb$bWbXb]bYb%bZb^b0b'bwcbdcdxcddedfd5bscmdgdyczcAc_b`bhdBctcucCcDcEcFcGcHcvcIcJcKcLcMc{bNcOcPcQcRcScTcac|bUcVcWcXcYcZc0c1c2c3c4c5c6cbc:c;c=c?c@c2 [c]c^c/cccdcecfcgchcicjckclcmcncoc_cpc7c8c9c!c#c$c%c'c(c)c*c+c,cRb2b-c.cR Q X Y / qcrc3b6b7b(b)b}b~b
226 h_graph = (<GraphDefinition>source)._h_graph 2nbod+b,b-b.b/bSb:bTb#b;b=b?b@bUb[bVb$bWbXb]bYb%bZb^b0b'bwcbdcdxcddedfd5bscmdgdyczcAc_b`bhdBctcucCcDcEcFcGcHcvcIcJcKcLcMc{bNcOcPcQcRcScTcac|bUcVcWcXcYcZc0c1c2c3c4c5c6cbc:c;c=c?c@c2 [c]c^c/cccdcecfcgchcicjckclcmcncoc_cpc7c8c9c!c#c$c%c'c(c)c*c+c,cRb2b-c.cR Q X Y / qcrc3b6b7b(b)b}b~b
227 else:
228 raise TypeError(
229 f"expected GraphBuilder or GraphDefinition, got {type(source).__name__}")
231 cdef cydriver.CUDA_GRAPH_INSTANTIATE_PARAMS params = cydriver.CUDA_GRAPH_INSTANTIATE_PARAMS( 2P nbF L M N O cbH : x z 3 y odG ; 5 = 6 4 b c d e f g h i j k l m n o p q A r s t u v B C D E +b,b-b.b/bSb:bTb#b;b=b?b@bUb[bVb$bWbXb]bYb%bZb^b0b'bwcbdcdxcddedfd5bscmdgdyczcAc_b`bhdBctcucCcDcEcFcGcHcvcIcJcKcLcMc{bNcOcPcQcRcScTc| ^ ` gbabdb} _ { hbbbeb+ 7 $ , 8 % 0 V W - 9 ' . ! ( ac|bUcVcWcXcYcZc0c1c2c3c4c5c6cbc:c;c=c?c@c2 [c]c^c/cccdcecfcgchcicjckclcmcncoc_cpc7c8c9c!c#c$c%c'c(c)c*c+c,cRb2b-c.cR Q X Y / qcrc3b6b7bw 1 (b)b) ? ] }b~bfb* @ K I ~ [
232 flags=0,
233 hUploadStream=<cydriver.CUstream>NULL, 2P nbF L M N O cbH : x z 3 y odG ; 5 = 6 4 b c d e f g h i j k l m n o p q A r s t u v B C D E +b,b-b.b/bSb:bTb#b;b=b?b@bUb[bVb$bWbXb]bYb%bZb^b0b'bwcbdcdxcddedfd5bscmdgdyczcAc_b`bhdBctcucCcDcEcFcGcHcvcIcJcKcLcMc{bNcOcPcQcRcScTc| ^ ` gbabdb} _ { hbbbeb+ 7 $ , 8 % 0 V W - 9 ' . ! ( ac|bUcVcWcXcYcZc0c1c2c3c4c5c6cbc:c;c=c?c@c2 [c]c^c/cccdcecfcgchcicjckclcmcncoc_cpc7c8c9c!c#c$c%c'c(c)c*c+c,cRb2b-c.cR Q X Y / qcrc3b6b7bw 1 (b)b) ? ] }b~bfb* @ K I ~ [
234 hErrNode_out=<cydriver.CUgraphNode>NULL, 2P nbF L M N O cbH : x z 3 y odG ; 5 = 6 4 b c d e f g h i j k l m n o p q A r s t u v B C D E +b,b-b.b/bSb:bTb#b;b=b?b@bUb[bVb$bWbXb]bYb%bZb^b0b'bwcbdcdxcddedfd5bscmdgdyczcAc_b`bhdBctcucCcDcEcFcGcHcvcIcJcKcLcMc{bNcOcPcQcRcScTc| ^ ` gbabdb} _ { hbbbeb+ 7 $ , 8 % 0 V W - 9 ' . ! ( ac|bUcVcWcXcYcZc0c1c2c3c4c5c6cbc:c;c=c?c@c2 [c]c^c/cccdcecfcgchcicjckclcmcncoc_cpc7c8c9c!c#c$c%c'c(c)c*c+c,cRb2b-c.cR Q X Y / qcrc3b6b7bw 1 (b)b) ? ] }b~bfb* @ K I ~ [
235 result_out=cydriver.CUgraphInstantiateResult.CUDA_GRAPH_INSTANTIATE_SUCCESS, 2P nbF L M N O cbH : x z 3 y odG ; 5 = 6 4 b c d e f g h i j k l m n o p q A r s t u v B C D E +b,b-b.b/bSb:bTb#b;b=b?b@bUb[bVb$bWbXb]bYb%bZb^b0b'bwcbdcdxcddedfd5bscmdgdyczcAc_b`bhdBctcucCcDcEcFcGcHcvcIcJcKcLcMc{bNcOcPcQcRcScTc| ^ ` gbabdb} _ { hbbbeb+ 7 $ , 8 % 0 V W - 9 ' . ! ( ac|bUcVcWcXcYcZc0c1c2c3c4c5c6cbc:c;c=c?c@c2 [c]c^c/cccdcecfcgchcicjckclcmcncoc_cpc7c8c9c!c#c$c%c'c(c)c*c+c,cRb2b-c.cR Q X Y / qcrc3b6b7bw 1 (b)b) ? ] }b~bfb* @ K I ~ [
236 )
237 if options: 2P nbF L M N O cbH : x z 3 y odG ; 5 = 6 4 b c d e f g h i j k l m n o p q A r s t u v B C D E +b,b-b.b/bSb:bTb#b;b=b?b@bUb[bVb$bWbXb]bYb%bZb^b0b'bwcbdcdxcddedfd5bscmdgdyczcAc_b`bhdBctcucCcDcEcFcGcHcvcIcJcKcLcMc{bNcOcPcQcRcScTc| ^ ` gbabdb} _ { hbbbeb+ 7 $ , 8 % 0 V W - 9 ' . ! ( ac|bUcVcWcXcYcZc0c1c2c3c4c5c6cbc:c;c=c?c@c2 [c]c^c/cccdcecfcgchcicjckclcmcncoc_cpc7c8c9c!c#c$c%c'c(c)c*c+c,cRb2b-c.cR Q X Y / qcrc3b6b7bw 1 (b)b) ? ] }b~bfb* @ K I ~ [
238 flags = 0 2odSbTb#bUbVb$bWbXbYb%bZb0b'bwcxcfd5b0 V W fb*
239 if options.auto_free_on_launch: 2odSbTb#bUbVb$bWbXbYb%bZb0b'bwcxcfd5b0 V W fb*
240 flags |= driver.CUgraphInstantiate_flags.CUDA_GRAPH_INSTANTIATE_FLAG_AUTO_FREE_ON_LAUNCH 2SbUbXbZbwcxc0 V W fb*
241 if options.upload_stream is not None: 2odSbTb#bUbVb$bWbXbYb%bZb0b'bwcxcfd5b0 V W fb*
242 flags |= driver.CUgraphInstantiate_flags.CUDA_GRAPH_INSTANTIATE_FLAG_UPLOAD 2od#b$b%b'b*
243 params.hUploadStream = as_cu(Stream_accept(options.upload_stream)._h_stream) 2od#b$b%b'b*
244 if options.device_launch: 2SbTb#bUbVb$bWbXbYb%bZb0b'bwcxcfd5b0 V W fb*
245 flags |= driver.CUgraphInstantiate_flags.CUDA_GRAPH_INSTANTIATE_FLAG_DEVICE_LAUNCH 2Wbfd5bfb*
246 if options.use_node_priority: 2SbTb#bUbVb$bWbXbYb%bZb0b'bwcxcfd5b0 V W fb*
247 flags |= driver.CUgraphInstantiate_flags.CUDA_GRAPH_INSTANTIATE_FLAG_USE_NODE_PRIORITY 2TbVbYb0bwcxc*
248 params.flags = flags 2SbTb#bUbVb$bWbXbYb%bZb0b'bwcxcfd5b0 V W fb*
250 # The exec is adopted only when result_out reports success, so the
251 # diagnostics below run before the handle is checked.
252 h_exec = create_graph_exec_handle(h_graph, ¶ms) 2P nbF L M N O cbH : x z 3 y G ; 5 = 6 4 b c d e f g h i j k l m n o p q A r s t u v B C D E +b,b-b.b/bSb:bTb#b;b=b?b@bUb[bVb$bWbXb]bYb%bZb^b0b'bwcbdcdxcddedfd5bscmdgdyczcAc_b`bhdBctcucCcDcEcFcGcHcvcIcJcKcLcMc{bNcOcPcQcRcScTc| ^ ` gbabdb} _ { hbbbeb+ 7 $ , 8 % 0 V W - 9 ' . ! ( ac|bUcVcWcXcYcZc0c1c2c3c4c5c6cbc:c;c=c?c@c2 [c]c^c/cccdcecfcgchcicjckclcmcncoc_cpc7c8c9c!c#c$c%c'c(c)c*c+c,cRb2b-c.cR Q X Y / qcrc3b6b7bw 1 (b)b) ? ] }b~bfb* @ K I ~ [
253 status = get_last_error() 2P nbF L M N O cbH : x z 3 y G ; 5 = 6 4 b c d e f g h i j k l m n o p q A r s t u v B C D E +b,b-b.b/bSb:bTb#b;b=b?b@bUb[bVb$bWbXb]bYb%bZb^b0b'bwcbdcdxcddedfd5bscmdgdyczcAc_b`bhdBctcucCcDcEcFcGcHcvcIcJcKcLcMc{bNcOcPcQcRcScTc| ^ ` gbabdb} _ { hbbbeb+ 7 $ , 8 % 0 V W - 9 ' . ! ( ac|bUcVcWcXcYcZc0c1c2c3c4c5c6cbc:c;c=c?c@c2 [c]c^c/cccdcecfcgchcicjckclcmcncoc_cpc7c8c9c!c#c$c%c'c(c)c*c+c,cRb2b-c.cR Q X Y / qcrc3b6b7bw 1 (b)b) ? ] }b~bfb* @ K I ~ [
254 if params.result_out == driver.CUgraphInstantiateResult.CUDA_GRAPH_INSTANTIATE_ERROR: 2P nbF L M N O cbH : x z 3 y G ; 5 = 6 4 b c d e f g h i j k l m n o p q A r s t u v B C D E +b,b-b.b/bSb:bTb#b;b=b?b@bUb[bVb$bWbXb]bYb%bZb^b0b'bwcbdcdxcddedfd5bscmdgdyczcAc_b`bhdBctcucCcDcEcFcGcHcvcIcJcKcLcMc{bNcOcPcQcRcScTc| ^ ` gbabdb} _ { hbbbeb+ 7 $ , 8 % 0 V W - 9 ' . ! ( ac|bUcVcWcXcYcZc0c1c2c3c4c5c6cbc:c;c=c?c@c2 [c]c^c/cccdcecfcgchcicjckclcmcncoc_cpc7c8c9c!c#c$c%c'c(c)c*c+c,cRb2b-c.cR Q X Y / qcrc3b6b7bw 1 (b)b) ? ] }b~bfb* @ K I ~ [
255 # HANDLE_RETURN raises CUDAError with the CUresult name and message (e.g. CUDA_ERROR_INVALID_VALUE)
256 # when status is not CUDA_SUCCESS.
257 HANDLE_RETURN(status) 2fb
258 raise RuntimeError(
259 "CUDA graph instantiation failed, but cuGraphInstantiateWithParams "
260 "returned CUDA_SUCCESS; no driver error details are available."
261 )
262 elif params.result_out == driver.CUgraphInstantiateResult.CUDA_GRAPH_INSTANTIATE_INVALID_STRUCTURE: 2P nbF L M N O cbH : x z 3 y G ; 5 = 6 4 b c d e f g h i j k l m n o p q A r s t u v B C D E +b,b-b.b/bSb:bTb#b;b=b?b@bUb[bVb$bWbXb]bYb%bZb^b0b'bwcbdcdxcddedfd5bscmdgdyczcAc_b`bhdBctcucCcDcEcFcGcHcvcIcJcKcLcMc{bNcOcPcQcRcScTc| ^ ` gbabdb} _ { hbbbeb+ 7 $ , 8 % 0 V W - 9 ' . ! ( ac|bUcVcWcXcYcZc0c1c2c3c4c5c6cbc:c;c=c?c@c2 [c]c^c/cccdcecfcgchcicjckclcmcncoc_cpc7c8c9c!c#c$c%c'c(c)c*c+c,cRb2b-c.cR Q X Y / qcrc3b6b7bw 1 (b)b) ? ] }b~b* @ K I ~ [
263 raise RuntimeError("Instantiation failed due to invalid structure, such as cycles.") 2sc
264 elif params.result_out == driver.CUgraphInstantiateResult.CUDA_GRAPH_INSTANTIATE_NODE_OPERATION_NOT_SUPPORTED: 2P nbF L M N O cbH : x z 3 y G ; 5 = 6 4 b c d e f g h i j k l m n o p q A r s t u v B C D E +b,b-b.b/bSb:bTb#b;b=b?b@bUb[bVb$bWbXb]bYb%bZb^b0b'bwcbdcdxcddedfd5bscmdgdyczcAc_b`bhdBctcucCcDcEcFcGcHcvcIcJcKcLcMc{bNcOcPcQcRcScTc| ^ ` gbabdb} _ { hbbbeb+ 7 $ , 8 % 0 V W - 9 ' . ! ( ac|bUcVcWcXcYcZc0c1c2c3c4c5c6cbc:c;c=c?c@c2 [c]c^c/cccdcecfcgchcicjckclcmcncoc_cpc7c8c9c!c#c$c%c'c(c)c*c+c,cRb2b-c.cR Q X Y / qcrc3b6b7bw 1 (b)b) ? ] }b~b* @ K I ~ [
265 raise RuntimeError( 2fd
266 "Instantiation for device launch failed because the graph contained an unsupported operation."
267 )
268 elif params.result_out == driver.CUgraphInstantiateResult.CUDA_GRAPH_INSTANTIATE_MULTIPLE_CTXS_NOT_SUPPORTED: 2P nbF L M N O cbH : x z 3 y G ; 5 = 6 4 b c d e f g h i j k l m n o p q A r s t u v B C D E +b,b-b.b/bSb:bTb#b;b=b?b@bUb[bVb$bWbXb]bYb%bZb^b0b'bwcbdcdxcdded5bscmdgdyczcAc_b`bhdBctcucCcDcEcFcGcHcvcIcJcKcLcMc{bNcOcPcQcRcScTc| ^ ` gbabdb} _ { hbbbeb+ 7 $ , 8 % 0 V W - 9 ' . ! ( ac|bUcVcWcXcYcZc0c1c2c3c4c5c6cbc:c;c=c?c@c2 [c]c^c/cccdcecfcgchcicjckclcmcncoc_cpc7c8c9c!c#c$c%c'c(c)c*c+c,cRb2b-c.cR Q X Y / qcrc3b6b7bw 1 (b)b) ? ] }b~b* @ K I ~ [
269 raise RuntimeError("Instantiation for device launch failed due to the nodes belonging to different contexts.") 25b
270 elif params.result_out == driver.CUgraphInstantiateResult.CUDA_GRAPH_INSTANTIATE_CONDITIONAL_HANDLE_UNUSED: 2P nbF L M N O cbH : x z 3 y G ; 5 = 6 4 b c d e f g h i j k l m n o p q A r s t u v B C D E +b,b-b.b/bSb:bTb#b;b=b?b@bUb[bVb$bWbXb]bYb%bZb^b0b'bwcbdcdxcdded5bscmdgdyczcAc_b`bhdBctcucCcDcEcFcGcHcvcIcJcKcLcMc{bNcOcPcQcRcScTc| ^ ` gbabdb} _ { hbbbeb+ 7 $ , 8 % 0 V W - 9 ' . ! ( ac|bUcVcWcXcYcZc0c1c2c3c4c5c6cbc:c;c=c?c@c2 [c]c^c/cccdcecfcgchcicjckclcmcncoc_cpc7c8c9c!c#c$c%c'c(c)c*c+c,cRb2b-c.cR Q X Y / qcrc3b6b7bw 1 (b)b) ? ] }b~b* @ K I ~ [
271 raise RuntimeError("One or more conditional handles are not associated with conditional builders.") 2md
272 elif params.result_out != driver.CUgraphInstantiateResult.CUDA_GRAPH_INSTANTIATE_SUCCESS: 2P nbF L M N O cbH : x z 3 y G ; 5 = 6 4 b c d e f g h i j k l m n o p q A r s t u v B C D E +b,b-b.b/bSb:bTb#b;b=b?b@bUb[bVb$bWbXb]bYb%bZb^b0b'bwcbdcdxcdded5bscgdyczcAc_b`bhdBctcucCcDcEcFcGcHcvcIcJcKcLcMc{bNcOcPcQcRcScTc| ^ ` gbabdb} _ { hbbbeb+ 7 $ , 8 % 0 V W - 9 ' . ! ( ac|bUcVcWcXcYcZc0c1c2c3c4c5c6cbc:c;c=c?c@c2 [c]c^c/cccdcecfcgchcicjckclcmcncoc_cpc7c8c9c!c#c$c%c'c(c)c*c+c,cRb2b-c.cR Q X Y / qcrc3b6b7bw 1 (b)b) ? ] }b~b* @ K I ~ [
273 raise RuntimeError(f"Graph instantiation failed with unexpected error code: {params.result_out}")
275 if as_cu(h_exec) == NULL: 2P nbF L M N O cbH : x z 3 y G ; 5 = 6 4 b c d e f g h i j k l m n o p q A r s t u v B C D E +b,b-b.b/bSb:bTb#b;b=b?b@bUb[bVb$bWbXb]bYb%bZb^b0b'bwcbdcdxcdded5bscgdyczcAc_b`bhdBctcucCcDcEcFcGcHcvcIcJcKcLcMc{bNcOcPcQcRcScTc| ^ ` gbabdb} _ { hbbbeb+ 7 $ , 8 % 0 V W - 9 ' . ! ( ac|bUcVcWcXcYcZc0c1c2c3c4c5c6cbc:c;c=c?c@c2 [c]c^c/cccdcecfcgchcicjckclcmcncoc_cpc7c8c9c!c#c$c%c'c(c)c*c+c,cRb2b-c.cR Q X Y / qcrc3b6b7bw 1 (b)b) ? ] }b~b* @ K I ~ [
276 HANDLE_RETURN(status)
277 return Graph._init(h_exec) 2P nbF L M N O cbH : x z 3 y G ; 5 = 6 4 b c d e f g h i j k l m n o p q A r s t u v B C D E +b,b-b.b/bSb:bTb#b;b=b?b@bUb[bVb$bWbXb]bYb%bZb^b0b'bwcbdcdxcdded5bscgdyczcAc_b`bhdBctcucCcDcEcFcGcHcvcIcJcKcLcMc{bNcOcPcQcRcScTc| ^ ` gbabdb} _ { hbbbeb+ 7 $ , 8 % 0 V W - 9 ' . ! ( ac|bUcVcWcXcYcZc0c1c2c3c4c5c6cbc:c;c=c?c@c2 [c]c^c/cccdcecfcgchcicjckclcmcncoc_cpc7c8c9c!c#c$c%c'c(c)c*c+c,cRb2b-c.cR Q X Y / qcrc3b6b7bw 1 (b)b) ? ] }b~b* @ K I ~ [
280# Distinguishes the three kinds of GraphBuilder, which differ in how they
281# begin/end stream capture and whether they own the resulting CUgraph.
282# Each kind progresses through _CaptureState as follows:
283#
284# PRIMARY: NOT_STARTED -> CAPTURING -> ENDED | INVALIDATED
285# FORKED: CAPTURING (never transitions; joined and closed)
286# CONDITIONAL_BODY: NOT_STARTED -> CAPTURING -> ENDED | INVALIDATED
287#
288cdef enum _BuilderKind:
289 # PRIMARY: The top-level builder created by Device or Stream. Owns the
290 # captured CUgraph via an owning GraphHandle. Progresses through all three
291 # capture states; responsible for ending capture if destroyed early. When
292 # the capture ends invalidated, the driver discards the graph and the
293 # handle is dropped without destroying it (see GB_end_capture).
294 PRIMARY = 0
295 # FORKED: Created by split(). Captures on a private stream forked from the
296 # primary. Starts in CAPTURING state and never transitions; the user joins
297 # it back to the primary via join(), which closes the builder. Must NOT
298 # call cuStreamEndCapture (the driver requires all forked streams to be
299 # joined first).
300 FORKED = 1
301 # CONDITIONAL_BODY: Created by if_then/if_else/switch/while_loop. Captures
302 # into a non-owned body graph via cuStreamBeginCaptureToGraph. The body
303 # graph's lifetime is tied to a parent graph. Progresses through all three
304 # capture states like PRIMARY.
305 CONDITIONAL_BODY = 2
308# Tracks the capture lifecycle of a GraphBuilder.
309cdef enum _CaptureState:
310 CAPTURE_NOT_STARTED = 0
311 CAPTURING = 1
312 CAPTURE_ENDED = 2 # Finished, valid handle
313 CLOSED = 3 # No valid handle
314 CAPTURE_INVALIDATED = 4 # Ended in error; the driver discarded the graph
317cdef class GraphBuilder:
318 """A graph under construction by stream capture.
320 A graph groups a set of CUDA kernels and other CUDA operations together and executes
321 them with a specified dependency tree. It speeds up the workflow by combining the
322 driver activities associated with CUDA kernel launches and CUDA API calls.
324 Directly creating a :obj:`~graph.GraphBuilder` is not supported due
325 to ambiguity. New graph builders should instead be created through a
326 :obj:`~_device.Device`, or a :obj:`~_stream.stream` object.
328 .. note::
330 Operations recorded during capture reference your memory but do not
331 take ownership of it. As with ordinary stream work, you must keep the
332 operands alive for as long as the completed graph may execute -- for
333 example, the :obj:`~_memory.Buffer` objects passed to :func:`~launch`
334 or :meth:`~_memory.Buffer.copy_to`. Host callbacks added with
335 :meth:`callback` are the exception: the callable (and any copied
336 ``user_data``) are retained for the graph's lifetime. This differs from
337 building a graph explicitly with :class:`~graph.GraphDefinition`, which
338 retains the operands it is given.
340 """
342 def __init__(self):
343 raise NotImplementedError( 2qd
344 "directly creating a GraphBuilder object can be ambiguous. Please either "
345 "call Device.create_graph_builder() or stream.create_graph_builder()"
346 )
348 def __dealloc__(self):
349 cdef cydriver.CUresult err = GB_end_capture(self) 2`c~cP Cbnb9bidpbobF lbibjbnd1bqdqbzbL M N O cbH : # x S z rbadQbZ sb3 mb4by *bDbG T ; ld5 = 6 kbpd4 b c d e f g h i j k l m n o p q A r s t u v B C D E J | ^ ` gbabdb!b8b} _ { hbbbeb+ 7 $ , 8 % 0 V W - 9 ' . ! ( 2 R Q X Y / w 1 ) ? ] Pbfb* a tbAbubvb@ {c|cK I ~ kdIbJbEbKbLbFbMbNbGbObHbwbxbybBb}c[
350 if err != cydriver.CUDA_SUCCESS: 2`c~cP Cbnb9bidpbobF lbibjbnd1bqdqbzbL M N O cbH : # x S z rbadQbZ sb3 mb4by *bDbG T ; ld5 = 6 kbpd4 b c d e f g h i j k l m n o p q A r s t u v B C D E J | ^ ` gbabdb!b8b} _ { hbbbeb+ 7 $ , 8 % 0 V W - 9 ' . ! ( 2 R Q X Y / w 1 ) ? ] Pbfb* a tbAbubvb@ {c|cK I ~ kdIbJbEbKbLbFbMbNbGbObHbwbxbybBb}c[
351 # Nothing can be raised from __dealloc__: report instead.
352 report_cuda_error( 29b
353 b"cuStreamEndCapture", err,
354 b"failed while releasing a GraphBuilder that was still building")
356 @staticmethod
357 def _init(Stream stream):
358 cdef GraphBuilder self = GraphBuilder.__new__(GraphBuilder) 2`c~cP Cbnb9bidpbobF lbibjbnd1bqbzbL M N O cbH : # x S z rbadQbZ sb3 mb4by *bDbG T U ; ld5 = 6 kbpd4 b c d e f g h i j k l m n o p q A r s t u v B C D E J | ^ ` gbabdb!b8b} _ { hbbbeb+ 7 $ , 8 % 0 V W - 9 ' . ! ( 2 R Q X Y / w 1 ) ? ] Pbfb* a tbAbubvb@ {c|cK I ~ kdIbJbEbKbLbFbMbNbGbObHbwbxbybBb}c[
359 # _h_graph set by begin_building
360 self._h_stream = stream._h_stream 2`c~cP Cbnb9bidpbobF lbibjbnd1bqbzbL M N O cbH : # x S z rbadQbZ sb3 mb4by *bDbG T U ; ld5 = 6 kbpd4 b c d e f g h i j k l m n o p q A r s t u v B C D E J | ^ ` gbabdb!b8b} _ { hbbbeb+ 7 $ , 8 % 0 V W - 9 ' . ! ( 2 R Q X Y / w 1 ) ? ] Pbfb* a tbAbubvb@ {c|cK I ~ kdIbJbEbKbLbFbMbNbGbObHbwbxbybBb}c[
361 self._kind = PRIMARY 2`c~cP Cbnb9bidpbobF lbibjbnd1bqbzbL M N O cbH : # x S z rbadQbZ sb3 mb4by *bDbG T U ; ld5 = 6 kbpd4 b c d e f g h i j k l m n o p q A r s t u v B C D E J | ^ ` gbabdb!b8b} _ { hbbbeb+ 7 $ , 8 % 0 V W - 9 ' . ! ( 2 R Q X Y / w 1 ) ? ] Pbfb* a tbAbubvb@ {c|cK I ~ kdIbJbEbKbLbFbMbNbGbObHbwbxbybBb}c[
362 self._state = CAPTURE_NOT_STARTED 2`c~cP Cbnb9bidpbobF lbibjbnd1bqbzbL M N O cbH : # x S z rbadQbZ sb3 mb4by *bDbG T U ; ld5 = 6 kbpd4 b c d e f g h i j k l m n o p q A r s t u v B C D E J | ^ ` gbabdb!b8b} _ { hbbbeb+ 7 $ , 8 % 0 V W - 9 ' . ! ( 2 R Q X Y / w 1 ) ? ] Pbfb* a tbAbubvb@ {c|cK I ~ kdIbJbEbKbLbFbMbNbGbObHbwbxbybBb}c[
363 self._stream = stream 2`c~cP Cbnb9bidpbobF lbibjbnd1bqbzbL M N O cbH : # x S z rbadQbZ sb3 mb4by *bDbG T U ; ld5 = 6 kbpd4 b c d e f g h i j k l m n o p q A r s t u v B C D E J | ^ ` gbabdb!b8b} _ { hbbbeb+ 7 $ , 8 % 0 V W - 9 ' . ! ( 2 R Q X Y / w 1 ) ? ] Pbfb* a tbAbubvb@ {c|cK I ~ kdIbJbEbKbLbFbMbNbGbObHbwbxbybBb}c[
364 return self 2`c~cP Cbnb9bidpbobF lbibjbnd1bqbzbL M N O cbH : # x S z rbadQbZ sb3 mb4by *bDbG T U ; ld5 = 6 kbpd4 b c d e f g h i j k l m n o p q A r s t u v B C D E J | ^ ` gbabdb!b8b} _ { hbbbeb+ 7 $ , 8 % 0 V W - 9 ' . ! ( 2 R Q X Y / w 1 ) ? ] Pbfb* a tbAbubvb@ {c|cK I ~ kdIbJbEbKbLbFbMbNbGbObHbwbxbybBb}c[
366 def close(self):
367 """Destroy the graph builder.
369 A builder that is still building ends its capture first. The builder
370 is closed even when that fails; the driver error is raised afterwards.
371 """
372 HANDLE_RETURN(GB_close(self)) 2`cCbnbF lbibjbqb: # rbZ mby G T U = 6 b c d e f g h i j k l m n o p q r s t u v a tbAbubvb{c|ckdIbJbEbKbLbFbMbNbGbObHbwbxbybBb}c
374 @property
375 def is_closed(self) -> bool:
376 """Whether this graph builder has been closed."""
377 return self._state == CLOSED 2CbnbT U }c
379 @property
380 def stream(self) -> Stream:
381 """Returns the stream associated with the graph builder."""
382 GB_check_open(self) 2`cF qbH : # x z rbZ sb3 mby *bG T U ; 5 = 6 4 b c d e f g h i j k l m n o p q A r s t u v B C D E gbabdbhbbbeb+ 7 $ , 8 % 0 V W - 9 ' . ! ( 2 R Q X Y / w 1 ) ? ] Pbfb* a tbubvb@ {c|c~ wbxbyb}c[
383 return self._stream 2`cF qbH : # x z rbZ sb3 mby *bG T U ; 5 = 6 4 b c d e f g h i j k l m n o p q A r s t u v B C D E gbabdbhbbbeb+ 7 $ , 8 % 0 V W - 9 ' . ! ( 2 R Q X Y / w 1 ) ? ] Pbfb* a tbubvb@ {c|c~ wbxbyb[
385 @property
386 def is_join_required(self) -> bool:
387 """Returns True if this graph builder must be joined before building is ended."""
388 return self._kind == FORKED 1F#ZyGTUbcdefghijklmnopqrstuva
390 @property
391 def graph_definition(self) -> GraphDefinition:
392 """The captured graph as an explicit :class:`~graph.GraphDefinition`.
394 .. versionadded:: 1.1.0
396 The returned :class:`~graph.GraphDefinition` is a view of the same
397 graph this builder is producing: nodes added through it appear in
398 subsequent :meth:`complete` and :meth:`debug_dot_print` calls, and
399 the view stays valid even after the builder is closed.
401 This lets you mix the capture and explicit APIs on a single graph,
402 for example to inspect what was captured, augment it with extra
403 nodes, or build a conditional body entirely with the explicit API.
405 Availability:
407 - **Primary builders** (created by :meth:`Device.create_graph_builder`
408 or :meth:`Stream.create_graph_builder`): only after
409 :meth:`end_building`.
411 - **Conditional-body builders** (returned by :meth:`if_then`,
412 :meth:`if_else`, :meth:`while_loop`, :meth:`switch`): both before
413 :meth:`begin_building` and after :meth:`end_building`. The body
414 graph already exists when the conditional is created, so you may
415 populate it through this view without ever calling
416 :meth:`begin_building` on the body builder.
418 - **Forked builders** (returned by :meth:`split`): never. Forked
419 builders share the primary builder's graph; access it through the
420 primary instead.
422 Returns
423 -------
424 GraphDefinition
425 A view of the graph being built.
427 Raises
428 ------
429 RuntimeError
430 If the builder is closed, forked, currently building, or (for
431 primary builders) has not started building yet. A
432 :class:`~graph.GraphDefinition` obtained before :meth:`close`
433 keeps working; only fresh access through this property is
434 rejected once the builder is closed.
435 """
436 GB_check_open(self) 2lbibjbx S z rbadQbZ sb3 mb4 J 2 R Q X Y /
437 if self._kind == FORKED: 2lbibjbx S z adQbZ sb3 mb4 J 2 R Q X Y /
438 raise RuntimeError( 1Z
439 "graph_definition is unavailable on forked graph builders; "
440 "access it through the primary builder instead."
441 )
442 elif self._state == CAPTURING: 2lbibjbx S z adQbsb3 mb4 J 2 R Q X Y /
443 raise RuntimeError( 2S Qb
444 "graph_definition is unavailable while capture is in "
445 "progress; call end_building() first."
446 )
447 GB_check_has_graph(self) 2lbibjbx z adsb3 mb4 J 2 R Q X Y /
448 if self._kind == PRIMARY: 2x z adsb3 mb4 J 2 R Q X Y /
449 if self._state == CAPTURE_NOT_STARTED: 2adsb3 mb4 J 2 R Q X Y /
450 raise RuntimeError( 2ad
451 "graph_definition is unavailable before begin_building() on "
452 "a primary builder; no graph has been created yet."
453 )
454 return GraphDefinition._from_handle(self._h_graph) 2x z sb3 mb4 J 2 R Q X Y /
456 def begin_building(self, mode: str | None = "relaxed") -> GraphBuilder:
457 """Begins the building process.
459 Build `mode` for controlling interaction with other API calls must be one of the following:
461 - `global` : Prohibit potentially unsafe operations across all streams in the process.
462 - `thread_local` : Prohibit potentially unsafe operations in streams created by the current thread.
463 - `relaxed` : The local thread is not prohibited from potentially unsafe operations.
465 Parameters
466 ----------
467 mode : str, optional
468 Build mode to control the interaction with other API calls that are porentially unsafe.
469 Default set to use relaxed.
471 """
472 GB_check_open(self) 2P Cb9bpbobF lbibjb1bqbzbL M N O cbH : # x S z rbQbZ sb3 mb4by *bDbG T U ; 5 = 6 kb4 b c d e f g h i j k l m n o p q A r s t u v B C D E J | ^ ` gbabdb!b8b} _ { hbbbeb+ 7 $ , 8 % 0 V W - 9 ' . ! ( 2 R Q X Y / w 1 ) ? ] Pbfb* a tbAbubvb@ K I ~ IbJbEbKbLbFbMbNbGbObHbwbxbybBb[
473 if self._state != CAPTURE_NOT_STARTED: 2P Cb9bpbobF lbibjb1bqbzbL M N O cbH : # x S z rbQbZ sb3 mb4by *bDbG T U ; 5 = 6 kb4 b c d e f g h i j k l m n o p q A r s t u v B C D E J | ^ ` gbabdb!b8b} _ { hbbbeb+ 7 $ , 8 % 0 V W - 9 ' . ! ( 2 R Q X Y / w 1 ) ? ] Pbfb* a tbAbubvb@ K I ~ IbJbEbKbLbFbMbNbGbObHbwbxbybBb[
474 if self._state == CAPTURING: 21b;
475 raise RuntimeError("Graph builder is already building.") 21b
476 else:
477 raise RuntimeError("Cannot resume building after building has ended.") 1;
478 cdef cydriver.CUstreamCaptureMode c_mode
479 if mode == "global": 2P Cb9bpbobF lbibjb1bqbzbL M N O cbH : # x S z rbQbZ sb3 mb4by *bDbG T U ; 5 = 6 kb4 b c d e f g h i j k l m n o p q A r s t u v B C D E J | ^ ` gbabdb!b8b} _ { hbbbeb+ 7 $ , 8 % 0 V W - 9 ' . ! ( 2 R Q X Y / w 1 ) ? ] Pbfb* a tbAbubvb@ K I ~ IbJbEbKbLbFbMbNbGbObHbwbxbybBb[
480 c_mode = cydriver.CU_STREAM_CAPTURE_MODE_GLOBAL 2lb| gb!b} hb+ , 0 - . PbIbJbKbLbMbNbOb
481 elif mode == "thread_local": 2P Cb9bpbobF ibjb1bqbzbL M N O cbH : # x S z rbQbZ sb3 mb4by *bDbG T U ; 5 = 6 kb4 b c d e f g h i j k l m n o p q A r s t u v B C D E J ^ ` abdb8b_ { bbeb7 $ 8 % V W 9 ' ! ( 2 R Q X Y / w 1 ) ? ] Pbfb* a tbAbubvb@ K I ~ EbFbGbHbwbxbybBb[
482 c_mode = cydriver.CU_STREAM_CAPTURE_MODE_THREAD_LOCAL 2jb` db8b{ eb$ % W ' ( PbEbFbGbHb
483 elif mode == "relaxed": 2P Cb9bpbobF ib1bqbzbL M N O cbH : # x S z rbQbZ sb3 mb4by *bDbG T U ; 5 = 6 kb4 b c d e f g h i j k l m n o p q A r s t u v B C D E J ^ ab_ bb7 8 V 9 ! 2 R Q X Y / w 1 ) ? ] Pbfb* a tbAbubvb@ K I ~ wbxbybBb[
484 c_mode = cydriver.CU_STREAM_CAPTURE_MODE_RELAXED 2P Cb9bpbobF ib1bqbzbL M N O cbH : # x S z rbQbZ sb3 mb4by *bDbG T U ; 5 = 6 kb4 b c d e f g h i j k l m n o p q A r s t u v B C D E J ^ ab_ bb7 8 V 9 ! 2 R Q X Y / w 1 ) ? ] Pbfb* a tbAbubvb@ K I ~ wbxbybBb[
485 else:
486 raise ValueError(f"Unsupported build mode: {mode}") 2Pb
488 cdef cydriver.CUstream c_stream = as_cu(self._h_stream) 2P Cb9bpbobF lbibjb1bqbzbL M N O cbH : # x S z rbQbZ sb3 mb4by *bDbG T U ; 5 = 6 kb4 b c d e f g h i j k l m n o p q A r s t u v B C D E J | ^ ` gbabdb!b8b} _ { hbbbeb+ 7 $ , 8 % 0 V W - 9 ' . ! ( 2 R Q X Y / w 1 ) ? ] Pbfb* a tbAbubvb@ K I ~ IbJbEbKbLbFbMbNbGbObHbwbxbybBb[
489 cdef cydriver.CUgraph c_graph
490 cdef cydriver.CUstreamCaptureStatus c_status
491 if self._kind == CONDITIONAL_BODY: 2P Cb9bpbobF lbibjb1bqbzbL M N O cbH : # x S z rbQbZ sb3 mb4by *bDbG T U ; 5 = 6 kb4 b c d e f g h i j k l m n o p q A r s t u v B C D E J | ^ ` gbabdb!b8b} _ { hbbbeb+ 7 $ , 8 % 0 V W - 9 ' . ! ( 2 R Q X Y / w 1 ) ? ] Pbfb* a tbAbubvb@ K I ~ IbJbEbKbLbFbMbNbGbObHbwbxbybBb[
492 c_graph = as_cu(self._h_graph) 1xSbcdefghijklmnopqArstuvBCDEwa
493 with nogil: 1xSbcdefghijklmnopqArstuvBCDEwa
494 HANDLE_RETURN(cydriver.cuStreamBeginCaptureToGraph( 1xSbcdefghijklmnopqArstuvBCDEwa
495 c_stream, c_graph, NULL, NULL, 0, c_mode))
496 self._state = CAPTURING 1xSbcdefghijklmnopqArstuvBCDEwa
497 else:
498 with nogil: 2P Cb9bpbobF lbibjb1bqbzbL M N O cbH : # x S z rbQbZ sb3 mb4by *bDbG T U ; 5 = 6 kb4 b c d e f g h i j k l m n o p q A r s t u v B C D E J | ^ ` gbabdb!b8b} _ { hbbbeb+ 7 $ , 8 % 0 V W - 9 ' . ! ( 2 R Q X Y / w 1 ) ? ] Pbfb* a tbAbubvb@ K I ~ IbJbEbKbLbFbMbNbGbObHbwbxbybBb[
499 HANDLE_RETURN(cydriver.cuStreamBeginCapture(c_stream, c_mode)) 2P Cb9bpbobF lbibjb1bqbzbL M N O cbH : # x S z rbQbZ sb3 mb4by *bDbG T U ; 5 = 6 kb4 b c d e f g h i j k l m n o p q A r s t u v B C D E J | ^ ` gbabdb!b8b} _ { hbbbeb+ 7 $ , 8 % 0 V W - 9 ' . ! ( 2 R Q X Y / w 1 ) ? ] Pbfb* a tbAbubvb@ K I ~ IbJbEbKbLbFbMbNbGbObHbwbxbybBb[
500 # Capture is active now; set CAPTURING before the calls below so a
501 # failure in _get_capture_info/create_graph_handle still lets
502 # cleanup end the capture rather than leaving the stream poisoned.
503 self._state = CAPTURING 2P Cb9bpbobF lbibjb1bqbzbL M N O cbH : # x S z rbQbZ sb3 mb4by *bDbG T U ; 5 = 6 kb4 b c d e f g h i j k l m n o p q A r s t u v B C D E J | ^ ` gbabdb!b8b} _ { hbbbeb+ 7 $ , 8 % 0 V W - 9 ' . ! ( 2 R Q X Y / w 1 ) ? ] Pbfb* a tbAbubvb@ K I ~ IbJbEbKbLbFbMbNbGbObHbwbxbybBb[
504 with nogil: 2P Cb9bpbobF lbibjb1bqbzbL M N O cbH : # x S z rbQbZ sb3 mb4by *bDbG T U ; 5 = 6 kb4 b c d e f g h i j k l m n o p q A r s t u v B C D E J | ^ ` gbabdb!b8b} _ { hbbbeb+ 7 $ , 8 % 0 V W - 9 ' . ! ( 2 R Q X Y / w 1 ) ? ] Pbfb* a tbAbubvb@ K I ~ IbJbEbKbLbFbMbNbGbObHbwbxbybBb[
505 # The driver rejects a NULL captureStatus_out, so pass a
506 # stack-local even though we only want the graph handle.
507 _get_capture_info(c_stream, &c_status, &c_graph) 2P Cb9bpbobF lbibjb1bqbzbL M N O cbH : # x S z rbQbZ sb3 mb4by *bDbG T U ; 5 = 6 kb4 b c d e f g h i j k l m n o p q A r s t u v B C D E J | ^ ` gbabdb!b8b} _ { hbbbeb+ 7 $ , 8 % 0 V W - 9 ' . ! ( 2 R Q X Y / w 1 ) ? ] Pbfb* a tbAbubvb@ K I ~ IbJbEbKbLbFbMbNbGbObHbwbxbybBb[
508 self._h_graph = create_graph_handle(c_graph) 2P Cb9bpbobF lbibjb1bqbzbL M N O cbH : # x S z rbQbZ sb3 mb4by *bDbG T U ; 5 = 6 kb4 b c d e f g h i j k l m n o p q A r s t u v B C D E J | ^ ` gbabdb!b8b} _ { hbbbeb+ 7 $ , 8 % 0 V W - 9 ' . ! ( 2 R Q X Y / w 1 ) ? ] Pbfb* a tbAbubvb@ K I ~ IbJbEbKbLbFbMbNbGbObHbwbxbybBb[
509 return self 2P Cb9bpbobF lbibjb1bqbzbL M N O cbH : # x S z rbQbZ sb3 mb4by *bDbG T U ; 5 = 6 kb4 b c d e f g h i j k l m n o p q A r s t u v B C D E J | ^ ` gbabdb!b8b} _ { hbbbeb+ 7 $ , 8 % 0 V W - 9 ' . ! ( 2 R Q X Y / w 1 ) ? ] Pbfb* a tbAbubvb@ K I ~ IbJbEbKbLbFbMbNbGbObHbwbxbybBb[
511 @property
512 def is_building(self) -> bool:
513 """Returns True if the graph builder is currently building."""
514 GB_check_open(self) 2oblbibjbqbH DbJ | ^ ` } _ {
515 cdef cydriver.CUstream c_stream = as_cu(self._h_stream) 2oblbibjbqbH DbJ | ^ ` } _ {
516 cdef cydriver.CUstreamCaptureStatus status
517 with nogil: 2oblbibjbqbH DbJ | ^ ` } _ {
518 _get_capture_info(c_stream, &status, NULL) 2oblbibjbqbH DbJ | ^ ` } _ {
519 if status == cydriver.CU_STREAM_CAPTURE_STATUS_NONE: 2oblbibjbqbH DbJ | ^ ` } _ {
520 return False 2oblbibjbDb
521 elif status == cydriver.CU_STREAM_CAPTURE_STATUS_ACTIVE:
522 return True 2qbH DbJ | ^ ` } _ {
523 elif status == cydriver.CU_STREAM_CAPTURE_STATUS_INVALIDATED:
524 raise RuntimeError( 2lbibjb
525 "Build process encountered an error and has been invalidated. "
526 "Call end_building() to end it; the error is raised there."
527 )
528 else:
529 raise NotImplementedError(f"Unsupported capture status type received: {status}")
531 def end_building(self) -> GraphBuilder:
532 """Ends the building process.
534 Raises
535 ------
536 RuntimeError
537 If the builder is not building. A forked builder is ended by
538 :meth:`join`, not by this method.
539 CUDAError
540 If the capture was invalidated, for example by a CUDA call that is
541 not permitted while capturing. The capture is ended and the
542 builder holds no graph afterwards, so :meth:`complete` and
543 :attr:`graph_definition` are unavailable. This applies to a
544 top-level builder. An invalidated capture of a conditional body
545 leaves the parent builder's graph invalid as well; ending it may
546 crash the process (see issue #2918).
547 """
548 GB_check_open(self) 2P pbobF lbibjbnd1bzbL M N O cbH : x S z rbQbZ sb3 mb4by DbG T U ; 5 = 6 kb4 b c d e f g h i j k l m n o p q A r s t u v B C D E J | ^ ` gbabdb!b8b} _ { hbbbeb+ 7 $ , 8 % 0 V W - 9 ' . ! ( 2 R Q X Y / w 1 ) ? ] Pbfb* a tbAbubvb@ K I ~ IbJbEbKbLbFbMbNbGbObHbwbxbybBb[
549 if self._state != CAPTURING: 2P pbobF lbibjbnd1bzbL M N O cbH : x S z rbQbZ sb3 mb4by DbG T U ; 5 = 6 kb4 b c d e f g h i j k l m n o p q A r s t u v B C D E J | ^ ` gbabdb!b8b} _ { hbbbeb+ 7 $ , 8 % 0 V W - 9 ' . ! ( 2 R Q X Y / w 1 ) ? ] Pbfb* a tbAbubvb@ K I ~ IbJbEbKbLbFbMbNbGbObHbwbxbybBb[
550 raise RuntimeError("Graph builder is not building.") 2nd
551 if self._kind == FORKED: 2P pbobF lbibjb1bzbL M N O cbH : x S z rbQbZ sb3 mb4by DbG T U ; 5 = 6 kb4 b c d e f g h i j k l m n o p q A r s t u v B C D E J | ^ ` gbabdb!b8b} _ { hbbbeb+ 7 $ , 8 % 0 V W - 9 ' . ! ( 2 R Q X Y / w 1 ) ? ] Pbfb* a tbAbubvb@ K I ~ IbJbEbKbLbFbMbNbGbObHbwbxbybBb[
552 raise RuntimeError( 1F
553 "A forked graph builder is ended by GraphBuilder.join(), not end_building().")
554 HANDLE_RETURN(GB_end_capture(self)) 2P pbobF lbibjb1bzbL M N O cbH : x S z rbQbZ sb3 mb4by DbG T U ; 5 = 6 kb4 b c d e f g h i j k l m n o p q A r s t u v B C D E J | ^ ` gbabdb!b8b} _ { hbbbeb+ 7 $ , 8 % 0 V W - 9 ' . ! ( 2 R Q X Y / w 1 ) ? ] Pbfb* a tbAbubvb@ K I ~ IbJbEbKbLbFbMbNbGbObHbwbxbybBb[
556 # TODO: Resolving https://github.com/NVIDIA/cuda-python/issues/617 would allow us to
557 # resume the build process after the first call to end_building()
558 return self 2P pbobF 1bzbL M N O cbH : x S z rbQbZ sb3 mb4by DbG ; 5 = 6 kb4 b c d e f g h i j k l m n o p q A r s t u v B C D E J | ^ ` gbabdb!b8b} _ { hbbbeb+ 7 $ , 8 % 0 V W - 9 ' . ! ( 2 R Q X Y / w 1 ) ? ] Pbfb* a tbAbubvb@ K I ~ IbJbEbKbLbFbMbNbGbObHbwbxbybBb[
560 def complete(self, options: GraphCompleteOptions | None = None) -> Graph:
561 """Completes the graph builder and returns the built :obj:`~graph.Graph` object.
563 Parameters
564 ----------
565 options : :obj:`~graph.GraphCompleteOptions`, optional
566 Customizable dataclass for the graph builder completion options.
568 Returns
569 -------
570 graph : :obj:`~graph.Graph`
571 The newly built graph.
573 """
574 GB_check_open(self) 2P F lbibjbL M N O cbH : # x z 3 y G ; 5 = 6 4 b c d e f g h i j k l m n o p q A r s t u v B C D E | ^ ` gbabdb} _ { hbbbeb+ 7 $ , 8 % 0 V W - 9 ' . ! ( w 1 ) ? ] fb* @ K I ~ [
575 GB_check_has_graph(self) 2P F lbibjbL M N O cbH : x z 3 y G ; 5 = 6 4 b c d e f g h i j k l m n o p q A r s t u v B C D E | ^ ` gbabdb} _ { hbbbeb+ 7 $ , 8 % 0 V W - 9 ' . ! ( w 1 ) ? ] fb* @ K I ~ [
576 if self._state != CAPTURE_ENDED: 2P F L M N O cbH : x z 3 y G ; 5 = 6 4 b c d e f g h i j k l m n o p q A r s t u v B C D E | ^ ` gbabdb} _ { hbbbeb+ 7 $ , 8 % 0 V W - 9 ' . ! ( w 1 ) ? ] fb* @ K I ~ [
577 raise RuntimeError("Graph has not finished building.") 2cb
579 return _instantiate_graph(self, options) 2P F L M N O cbH : x z 3 y G ; 5 = 6 4 b c d e f g h i j k l m n o p q A r s t u v B C D E | ^ ` gbabdb} _ { hbbbeb+ 7 $ , 8 % 0 V W - 9 ' . ! ( w 1 ) ? ] fb* @ K I ~ [
581 def debug_dot_print(self, path: str, options: GraphDebugPrintOptions | None = None) -> None:
582 """Generates a DOT debug file for the graph builder.
584 Parameters
585 ----------
586 path : str
587 File path to use for writting debug DOT output
588 options : :obj:`~graph.GraphDebugPrintOptions`, optional
589 Customizable dataclass for the debug print options.
591 """
592 GB_check_open(self) 2pba
593 GB_check_has_graph(self) 2pba
594 if self._state != CAPTURE_ENDED: 2pba
595 raise RuntimeError("Graph has not finished building.") 2pb
596 cdef unsigned int c_flags = options._to_flags() if options else 0 2pba
597 cdef cydriver.CUgraph c_graph = as_cu(self._h_graph) 2pba
598 cdef bytes b_path = path.encode('utf-8') 2pba
599 cdef const char* c_path = b_path 2pba
600 with nogil: 2pba
601 HANDLE_RETURN(cydriver.cuGraphDebugDotPrint(c_graph, c_path, c_flags)) 2pba
603 def split(self, count: int) -> tuple[GraphBuilder, ...]:
604 """Splits the original graph builder into multiple graph builders.
606 The new builders inherit work dependencies from the original builder.
607 The original builder is reused for the split and is returned first in the tuple.
609 Parameters
610 ----------
611 count : int
612 The number of graph builders to split the graph builder into.
614 Returns
615 -------
616 graph_builders : tuple[:obj:`~graph.GraphBuilder`, ...]
617 A tuple of split graph builders. The first graph builder in the tuple
618 is always the original graph builder.
620 """
621 if count < 2: 2F # Z y G T U ldb c d e f g h i j k l m n o p q r s t u v a
622 raise ValueError(f"Invalid split count: expecting >= 2, got {count}") 1y
623 GB_check_open(self) 2F # Z y G T U ldb c d e f g h i j k l m n o p q r s t u v a
624 if self._state != CAPTURING: 2F # Z y G T U ldb c d e f g h i j k l m n o p q r s t u v a
625 raise RuntimeError("Graph builder must be building before it can be split.") 2ld
627 event = self._stream.record() 1F#ZyGTUbcdefghijklmnopqrstuva
628 result = [self] 1F#ZyGTUbcdefghijklmnopqrstuva
629 for _ in range(count - 1): 1F#ZyGTUbcdefghijklmnopqrstuva
630 stream = self._stream.device.create_stream() 1F#ZyGTUbcdefghijklmnopqrstuva
631 stream.wait(event) 1F#ZyGTUbcdefghijklmnopqrstuva
632 result.append(GB_init_forked(stream, self._h_graph)) 1F#ZyGTUbcdefghijklmnopqrstuva
633 event.close() 1F#ZyGTUbcdefghijklmnopqrstuva
634 return tuple(result) 1F#ZyGTUbcdefghijklmnopqrstuva
636 @staticmethod
637 def join(*graph_builders: GraphBuilder) -> GraphBuilder:
638 """Joins multiple graph builders into a single graph builder.
640 The returned builder inherits work dependencies from the provided builders.
641 If joining fails partway, the builders that were not joined are closed
642 before the error propagates, so none is left capturing. A driver error
643 from one of those closes is attached to the propagating error as a note.
645 Parameters
646 ----------
647 *graph_builders : :obj:`~graph.GraphBuilder`
648 The graph builders to join.
650 Returns
651 -------
652 graph_builder : :obj:`~graph.GraphBuilder`
653 The newly joined graph builder.
655 """
656 if any(not isinstance(builder, GraphBuilder) for builder in graph_builders): 2nbF # Z y G T U pdb c d e f g h i j k l m n o p q r s t u v a
657 raise TypeError("All arguments must be GraphBuilder instances") 2pd
658 if len(graph_builders) < 2: 2nbF # Z y G T U b c d e f g h i j k l m n o p q r s t u v a
659 raise ValueError("Must join with at least two graph builders") 1y
660 for builder in graph_builders: 2nbF # Z y G T U b c d e f g h i j k l m n o p q r s t u v a
661 GB_check_open(builder) 2nbF # Z y G T U b c d e f g h i j k l m n o p q r s t u v a
663 # Discover the root builder others should join
664 root_idx = 0 1F#ZyGTUbcdefghijklmnopqrstuva
665 for i, builder in enumerate(graph_builders): 1F#ZyGTUbcdefghijklmnopqrstuva
666 if not builder.is_join_required: 1F#ZyGTUbcdefghijklmnopqrstuva
667 root_idx = i 1F#ZyGTUbcdefghijklmnopqrstuva
668 break 1F#ZyGTUbcdefghijklmnopqrstuva
670 # Join all onto the root builder
671 root_bdr = graph_builders[root_idx] 1F#ZyGTUbcdefghijklmnopqrstuva
672 cdef GraphBuilder unjoined
673 try: 1F#ZyGTUbcdefghijklmnopqrstuva
674 for idx, builder in enumerate(graph_builders): 1F#ZyGTUbcdefghijklmnopqrstuva
675 if idx == root_idx: 1F#ZyGTUbcdefghijklmnopqrstuva
676 continue 1F#ZyGTUbcdefghijklmnopqrstuva
677 root_bdr.stream.wait(builder.stream) 1F#ZyGTUbcdefghijklmnopqrstuva
678 builder.close() 1F#ZyGbcdefghijklmnopqrstuva
679 except BaseException: 1TU
680 # A fork left open mid-capture crashed the interpreter when it was
681 # collected later (#2776). Close whatever the loop did not reach. A
682 # close that fails (a builder from another capture, which the
683 # refused wait has invalidated) must neither stop this sweep nor
684 # replace the error being raised, so its status is attached to
685 # that error as a note (error handling policy).
686 for idx, builder in enumerate(graph_builders): 1TU
687 if idx != root_idx and not builder.is_closed: 1TU
688 unjoined = builder 1TU
689 attach_rollback_failure( 1TU
690 b"cuStreamEndCapture", GB_close(unjoined),
691 b"failed while closing a graph builder that GraphBuilder.join did not join")
692 raise 1TU
694 return root_bdr 1F#ZyGbcdefghijklmnopqrstuva
696 def __cuda_stream__(self) -> tuple[int, int]:
697 """Return an instance of a __cuda_stream__ protocol."""
698 GB_check_open(self) 2`c
699 return self.stream.__cuda_stream__() 2`c
701 def _get_conditional_context(self) -> driver.CUcontext:
702 return self._stream.context.handle 2x S z kbb c d e f g h i j k l m n o p q A r s t u v B C D E w a
704 def create_condition(self, default_value: int | None = None) -> GraphCondition:
705 """Create a condition variable for use with conditional nodes.
707 The returned :class:`GraphCondition` object is passed to conditional-node
708 builder methods (:meth:`if_then`, :meth:`if_else`, :meth:`while_loop`,
709 :meth:`switch`). Its value is controlled at runtime by device code via
710 ``cudaGraphSetConditional``.
712 Parameters
713 ----------
714 default_value : int, optional
715 The default value to assign to the condition. If None, no
716 default is assigned.
718 Returns
719 -------
720 GraphCondition
721 A condition variable for controlling conditional execution.
722 """
723 GB_check_open(self) 2idx S z kbb c d e f g h i j k l m n o p q A r s t u v B C D E w a
724 if cy_driver_version() < (12, 3, 0): 2idx S z kbb c d e f g h i j k l m n o p q A r s t u v B C D E w a
725 raise RuntimeError(f"Driver version {'.'.join(map(str, cy_driver_version()))} does not support conditional handles")
726 if default_value is not None: 2idx S z kbb c d e f g h i j k l m n o p q A r s t u v B C D E w a
727 flags = driver.CU_GRAPH_COND_ASSIGN_DEFAULT 2x S z kbB C D E w
728 else:
729 default_value = 0 2idb c d e f g h i j k l m n o p q A r s t u v a
730 flags = 0 2idb c d e f g h i j k l m n o p q A r s t u v a
732 status, _, graph, *_, _ = handle_return(driver.cuStreamGetCaptureInfo(self._stream.handle)) 2idx S z kbb c d e f g h i j k l m n o p q A r s t u v B C D E w a
733 if status != driver.CUstreamCaptureStatus.CU_STREAM_CAPTURE_STATUS_ACTIVE: 2idx S z kbb c d e f g h i j k l m n o p q A r s t u v B C D E w a
734 raise RuntimeError("Cannot create a condition when graph is not being built") 2id
736 raw_handle = handle_return( 2x S z kbb c d e f g h i j k l m n o p q A r s t u v B C D E w a
737 driver.cuGraphConditionalHandleCreate(graph, self._get_conditional_context(), default_value, flags) 2x S z kbb c d e f g h i j k l m n o p q A r s t u v B C D E w a
738 )
739 return GraphCondition._from_handle(<cydriver.CUgraphConditionalHandle><intptr_t>int(raw_handle)) 2x S z kbb c d e f g h i j k l m n o p q A r s t u v B C D E w a
741 def if_then(self, GraphCondition condition not None: GraphCondition) -> GraphBuilder:
742 """Adds an if condition branch and returns a new graph builder for it.
744 The resulting if graph will only execute the branch if the
745 condition evaluates to true at runtime.
747 The new builder inherits work dependencies from the original builder.
749 Parameters
750 ----------
751 condition : :class:`~graph.GraphCondition`
752 The condition variable from :meth:`create_condition` controlling
753 whether the branch executes.
755 Returns
756 -------
757 graph_builder : :obj:`~graph.GraphBuilder`
758 The newly created conditional graph builder.
760 """
761 GB_check_open(self) 2x S z kbb c d e f g h i A r a
762 if cy_driver_version() < (12, 3, 0): 2x S z kbb c d e f g h i A r a
763 raise RuntimeError(f"Driver version {'.'.join(map(str, cy_driver_version()))} does not support conditional if")
764 if not isinstance(condition, GraphCondition): 2x S z kbb c d e f g h i A r a
765 raise TypeError(
766 f"condition must be a GraphCondition object (from "
767 f"GraphBuilder.create_condition()), got {type(condition).__name__}")
768 node_params = driver.CUgraphNodeParams() 2x S z kbb c d e f g h i A r a
769 node_params.type = driver.CUgraphNodeType.CU_GRAPH_NODE_TYPE_CONDITIONAL 2x S z kbb c d e f g h i A r a
770 node_params.conditional.handle = condition.handle 2x S z kbb c d e f g h i A r a
771 node_params.conditional.type = driver.CUgraphConditionalNodeType.CU_GRAPH_COND_TYPE_IF 2x S z kbb c d e f g h i A r a
772 node_params.conditional.size = 1 2x S z kbb c d e f g h i A r a
773 node_params.conditional.ctx = self._get_conditional_context() 2x S z kbb c d e f g h i A r a
774 return GB_cond_with_params(self, node_params)[0] 2x S z kbb c d e f g h i A r a
776 def if_else(self, GraphCondition condition not None: GraphCondition) -> tuple[GraphBuilder, GraphBuilder]:
777 """Adds an if-else condition branch and returns new graph builders for both branches.
779 The resulting if graph will execute the branch if the condition
780 evaluates to true at runtime, otherwise the else branch will execute.
782 The new builders inherit work dependencies from the original builder.
784 Parameters
785 ----------
786 condition : :class:`~graph.GraphCondition`
787 The condition variable from :meth:`create_condition` controlling
788 which branch executes.
790 Returns
791 -------
792 graph_builders : tuple[:obj:`~graph.GraphBuilder`, :obj:`~graph.GraphBuilder`]
793 A tuple of two new graph builders, one for the if branch and one for the else branch.
795 """
796 GB_check_open(self) 1jklmnopq
797 if cy_driver_version() < (12, 8, 0): 1jklmnopq
798 raise RuntimeError(f"Driver version {'.'.join(map(str, cy_driver_version()))} does not support conditional if-else")
799 if not isinstance(condition, GraphCondition): 1jklmnopq
800 raise TypeError(
801 f"condition must be a GraphCondition object (from "
802 f"GraphBuilder.create_condition()), got {type(condition).__name__}")
803 node_params = driver.CUgraphNodeParams() 1jklmnopq
804 node_params.type = driver.CUgraphNodeType.CU_GRAPH_NODE_TYPE_CONDITIONAL 1jklmnopq
805 node_params.conditional.handle = condition.handle 1jklmnopq
806 node_params.conditional.type = driver.CUgraphConditionalNodeType.CU_GRAPH_COND_TYPE_IF 1jklmnopq
807 node_params.conditional.size = 2 1jklmnopq
808 node_params.conditional.ctx = self._get_conditional_context() 1jklmnopq
809 return GB_cond_with_params(self, node_params) 1jklmnopq
811 def switch(self, GraphCondition condition not None: GraphCondition, count: int) -> tuple[GraphBuilder, ...]:
812 """Adds a switch condition branch and returns new graph builders for all cases.
814 The resulting switch graph will execute the branch whose case index
815 matches the value of the condition at runtime. If no match is found, no
816 branch will be executed.
818 The new builders inherit work dependencies from the original builder.
820 Parameters
821 ----------
822 condition : :class:`~graph.GraphCondition`
823 The condition variable from :meth:`create_condition` selecting
824 which case executes.
825 count : int
826 The number of cases to add to the switch conditional.
828 Returns
829 -------
830 graph_builders : tuple[:obj:`~graph.GraphBuilder`, ...]
831 A tuple of new graph builders, one for each branch.
833 """
834 GB_check_open(self) 1stuvw
835 if cy_driver_version() < (12, 8, 0): 1stuvw
836 raise RuntimeError(f"Driver version {'.'.join(map(str, cy_driver_version()))} does not support conditional switch")
837 if not isinstance(condition, GraphCondition): 1stuvw
838 raise TypeError(
839 f"condition must be a GraphCondition object (from "
840 f"GraphBuilder.create_condition()), got {type(condition).__name__}")
841 node_params = driver.CUgraphNodeParams() 1stuvw
842 node_params.type = driver.CUgraphNodeType.CU_GRAPH_NODE_TYPE_CONDITIONAL 1stuvw
843 node_params.conditional.handle = condition.handle 1stuvw
844 node_params.conditional.type = driver.CUgraphConditionalNodeType.CU_GRAPH_COND_TYPE_SWITCH 1stuvw
845 node_params.conditional.size = count 1stuvw
846 node_params.conditional.ctx = self._get_conditional_context() 1stuvw
847 return GB_cond_with_params(self, node_params) 1stuvw
849 def while_loop(self, GraphCondition condition not None: GraphCondition) -> GraphBuilder:
850 """Adds a while loop and returns a new graph builder for it.
852 The resulting while loop graph will execute the branch repeatedly at runtime
853 until the condition evaluates to false.
855 The new builder inherits work dependencies from the original builder.
857 Parameters
858 ----------
859 condition : :class:`~graph.GraphCondition`
860 The condition variable from :meth:`create_condition` controlling
861 loop continuation.
863 Returns
864 -------
865 graph_builder : :obj:`~graph.GraphBuilder`
866 The newly created while loop graph builder.
868 """
869 GB_check_open(self) 1BCDE
870 if cy_driver_version() < (12, 3, 0): 1BCDE
871 raise RuntimeError(f"Driver version {'.'.join(map(str, cy_driver_version()))} does not support conditional while loop")
872 if not isinstance(condition, GraphCondition): 1BCDE
873 raise TypeError(
874 f"condition must be a GraphCondition object (from "
875 f"GraphBuilder.create_condition()), got {type(condition).__name__}")
876 node_params = driver.CUgraphNodeParams() 1BCDE
877 node_params.type = driver.CUgraphNodeType.CU_GRAPH_NODE_TYPE_CONDITIONAL 1BCDE
878 node_params.conditional.handle = condition.handle 1BCDE
879 node_params.conditional.type = driver.CUgraphConditionalNodeType.CU_GRAPH_COND_TYPE_WHILE 1BCDE
880 node_params.conditional.size = 1 1BCDE
881 node_params.conditional.ctx = self._get_conditional_context() 1BCDE
882 return GB_cond_with_params(self, node_params)[0] 1BCDE
884 def embed(self, GraphBuilder child):
885 """Embed a previously-built :obj:`~graph.GraphBuilder` as a child node.
887 Parameters
888 ----------
889 child : :obj:`~graph.GraphBuilder`
890 The child graph builder. Must have finished building.
891 """
892 GB_check_open(self) 2obH J
893 GB_check_open(child) 2obH J
894 GB_check_has_graph(child) 2obH J
895 if child._state != CAPTURE_ENDED: 2obH J
896 raise ValueError("Child graph has not finished building.") 2ob
898 if not self.is_building: 2obH J
899 raise ValueError("Parent graph is not being built.") 2ob
901 stream_handle = self._stream.handle 1HJ
902 _, _, graph_out, *deps_info_out, num_dependencies_out = handle_return( 1HJ
903 driver.cuStreamGetCaptureInfo(stream_handle) 1HJ
904 )
906 # See https://github.com/NVIDIA/cuda-python/pull/879#issuecomment-3211054159
907 # for rationale
908 dependencies_out = deps_info_out[0] 1HJ
909 new_node = handle_return( 1HJ
910 driver.cuGraphAddChildGraphNode( 1HJ
911 graph_out, dependencies_out,
912 num_dependencies_out, as_py(child._h_graph) 1HJ
913 )
914 )
915 cdef cydriver.CUgraphNode c_new_node = (
916 <cydriver.CUgraphNode><intptr_t>int(new_node) 1HJ
917 )
918 cdef cydriver.CUgraph embedded_graph = NULL 1HJ
919 cdef cydriver.CUresult rollback_status
920 cdef GraphHandle h_embedded
921 try: 1HJ
922 with nogil: 1HJ
923 HANDLE_RETURN(cydriver.cuGraphChildGraphNodeGetGraph( 1HJ
924 c_new_node, &embedded_graph))
925 h_embedded = create_child_graph_handle( 1HJ
926 embedded_graph, self._h_graph, c_new_node)
927 HANDLE_RETURN(graph_clone_attachments( 1HJ
928 h_embedded, child._h_graph))
929 except:
930 with nogil:
931 rollback_status = cydriver.cuGraphDestroyNode(c_new_node)
932 if rollback_status == cydriver.CUDA_SUCCESS:
933 invalidate_child_graph_state(
934 self._h_graph, c_new_node)
935 else:
936 # The original exception propagates with the failed rollback
937 # attached as a note (error handling policy).
938 attach_rollback_failure(
939 b"cuGraphDestroyNode", rollback_status,
940 b"failed while rolling back a child graph node; the node remains in the graph")
941 raise
943 deps_info_update = [[new_node]] + [None] * (len(deps_info_out) - 1) 1HJ
944 handle_return( 1HJ
945 driver.cuStreamUpdateCaptureDependencies( 1HJ
946 stream_handle, 1HJ
947 *deps_info_update, # dependencies, edgeData
948 1,
949 driver.CUstreamUpdateCaptureDependencies_flags.CU_STREAM_SET_CAPTURE_DEPENDENCIES, 1HJ
950 )
951 )
953 def callback(self, fn: Callable[..., object], *, user_data: int | BufferProtocol | None = None) -> None:
954 """Add a host callback to the graph during stream capture.
956 The callback runs on the host CPU when the graph reaches this point
957 in execution. Two modes are supported:
959 - **Python callable**: Pass any callable. The GIL is acquired
960 automatically. The callable must take no arguments; use closures
961 or ``functools.partial`` to bind state.
962 - **ctypes function pointer**: The function receives a single
963 ``void*`` argument (the ``user_data``), and the caller must keep
964 the ctypes wrapper alive for the lifetime of the graph. Its
965 declared prototype must match the driver's ``CUhostFn``
966 (``void (*)(void*)``): ``ctypes.CFUNCTYPE(None, ctypes.c_void_p)``,
967 or ``ctypes.WINFUNCTYPE(None, ctypes.c_void_p)`` on Windows.
969 .. warning::
971 Callbacks must not call CUDA API functions. Doing so may
972 deadlock or corrupt driver state.
974 Use caution when a Python callback retains an object that owns a
975 graph. Any reference cycle involving the callback and a graph that
976 retains it cannot be broken by Python's cyclic garbage collector.
977 Use a weak reference to break such cycles.
979 Parameters
980 ----------
981 fn : callable or ctypes function pointer
982 The callback function.
983 user_data : int or bytes-like, optional
984 Only for ctypes function pointers. If ``int``, passed as a raw
985 pointer (caller manages lifetime). If bytes-like, the data is
986 copied and its lifetime is tied to the graph.
988 Raises
989 ------
990 TypeError
991 If ``fn`` is a ctypes function pointer whose declared prototype
992 does not match ``CUhostFn``.
993 ValueError
994 If ``user_data`` is given for a Python callable, or if ``fn`` is
995 :data:`HOST_COPY_FN` and ``user_data`` is not a packed record.
996 """
997 GB_callback(self, fn, user_data, False) 2~czbL M N O J K I
1000cdef inline void GB_callback(
1001 GraphBuilder gb, object fn, object user_data,
1002 bint fail_tail_discovery_for_testing) except *:
1003 GB_check_open(gb) 2~cP zbL M N O J K I
1004 cdef Stream stream = gb._stream 2~cP zbL M N O J K I
1005 cdef cydriver.CUstream c_stream = as_cu(stream._h_stream) 2~cP zbL M N O J K I
1006 cdef cydriver.CUstreamCaptureStatus capture_status
1008 with nogil: 2~cP zbL M N O J K I
1009 _get_capture_info(c_stream, &capture_status, NULL) 2~cP zbL M N O J K I
1011 if capture_status != cydriver.CU_STREAM_CAPTURE_STATUS_ACTIVE: 2~cP zbL M N O J K I
1012 raise RuntimeError("Cannot add callback when graph is not being built") 2~c
1014 cdef cydriver.CUhostFn c_fn
1015 cdef void* c_user_data = NULL 2P zbL M N O J K I
1016 cdef OpaqueHandle fn_owner, data_owner
1017 cdef PreparedAttachment prepared
1018 _resolve_host_callback( 2P zbL M N O J K I
1019 fn, user_data, &c_fn, &c_user_data, &fn_owner, &data_owner)
1020 HANDLE_RETURN(graph_prepare_attachment( 1PLMNOJKI
1021 gb._h_graph, fn_owner, data_owner, &prepared))
1023 with nogil: 1PLMNOJKI
1024 HANDLE_RETURN(cydriver.cuLaunchHostFunc(c_stream, c_fn, c_user_data)) 1PLMNOJKI
1026 # Capturing the host function added a node to the graph; it is now the
1027 # stream's sole capture dependency. Attach the callback's owners to it.
1028 cdef cydriver.CUgraphNode host_node
1029 cdef cydriver.CUresult commit_status
1030 try: 1PLMNOJKI
1031 if fail_tail_discovery_for_testing: 1PLMNOJKI
1032 raise RuntimeError("forced capture tail discovery failure") 1P
1033 host_node = _capture_tail_node(c_stream) 1LMNOJKI
1034 except BaseException as orig_exc: 1P
1035 # CUDA added the callback, but its node cannot be identified.
1036 # Retain its owners anonymously to prevent dangling pointers.
1037 commit_status = graph_commit_attachment(prepared, NULL) 1P
1038 try: 1P
1039 HANDLE_RETURN(commit_status) 1P
1040 except Exception as commit_exc:
1041 raise commit_exc from orig_exc
1042 raise 1P
1043 HANDLE_RETURN(graph_commit_attachment(prepared, host_node)) 1LMNOJKI
1046def _capture_callback_with_tail_failure_for_testing(
1047 GraphBuilder gb, fn, *, user_data=None):
1048 """Exercise anonymous attachment retention after node discovery fails."""
1049 GB_callback(gb, fn, user_data, True) 1P
1052cdef inline int GB_check_open(GraphBuilder gb) except -1:
1053 """Reject operations on a builder that has been closed.
1055 A CLOSED builder has reset its stream and graph handles, so any method
1056 that dereferences them would read a null handle (or, for the cached
1057 Stream, a None typed as cdef Stream). Guarding here yields a clear error
1058 instead.
1059 """
1060 if gb._state == CLOSED: 2`c~cP Cbnb9bidpbobF lbibjbnd1bqbzbL M N O cbH : # x S z rbadQbZ sb3 mb4by *bDbG T U ; ld5 = 6 kb4 b c d e f g h i j k l m n o p q A r s t u v B C D E J | ^ ` gbabdb!b8b} _ { hbbbeb+ 7 $ , 8 % 0 V W - 9 ' . ! ( 2 R Q X Y / w 1 ) ? ] Pbfb* a tbAbubvb@ {c|cK I ~ IbJbEbKbLbFbMbNbGbObHbwbxbybBb}c[
1061 raise RuntimeError("GraphBuilder has been closed") 2`cnbqb# rb6 }c
1062 return 0 2`c~cP Cb9bidpbobF lbibjbnd1bqbzbL M N O cbH : # x S z rbadQbZ sb3 mb4by *bDbG T U ; ld5 = 6 kb4 b c d e f g h i j k l m n o p q A r s t u v B C D E J | ^ ` gbabdb!b8b} _ { hbbbeb+ 7 $ , 8 % 0 V W - 9 ' . ! ( 2 R Q X Y / w 1 ) ? ] Pbfb* a tbAbubvb@ {c|cK I ~ IbJbEbKbLbFbMbNbGbObHbwbxbybBb[
1065cdef inline int GB_check_has_graph(GraphBuilder gb) except -1:
1066 """Reject graph access on a builder whose capture ended invalidated.
1068 The driver discards the capture graph when it ends an invalidated
1069 capture, so such a builder has nothing to complete, print, embed, or
1070 update from.
1071 """
1072 if gb._state == CAPTURE_INVALIDATED: 2P pbobF lbibjbL M N O cbH : x z adsb3 mby G ; 5 = 6 4 b c d e f g h i j k l m n o p q A r s t u v B C D E J | ^ ` gbabdb} _ { hbbbeb+ 7 $ , 8 % 0 V W - 9 ' . ! ( 2 R Q X Y / w 1 ) ? ] fb* a @ K I ~ [
1073 raise RuntimeError( 2lbibjb
1074 "The capture was invalidated and the driver discarded its graph; "
1075 "this graph builder has no graph.")
1076 return 0 2P pbobF L M N O cbH : x z adsb3 mby G ; 5 = 6 4 b c d e f g h i j k l m n o p q A r s t u v B C D E J | ^ ` gbabdb} _ { hbbbeb+ 7 $ , 8 % 0 V W - 9 ' . ! ( 2 R Q X Y / w 1 ) ? ] fb* a @ K I ~ [
1079cdef inline cydriver.CUresult GB_end_capture(GraphBuilder gb) noexcept:
1080 """End the capture if this builder owns a live one, and settle ownership
1081 of the graph. Returns the driver status; the caller raises or reports it.
1083 Only a CAPTURING PRIMARY or CONDITIONAL_BODY builder owns the live
1084 capture. A FORKED builder must not call cuStreamEndCapture: the driver
1085 requires forked streams to be joined first.
1087 On success the driver hands back the graph this builder already holds.
1088 When the capture was invalidated (or left unjoined), the driver destroys
1089 the capture graph and returns NULL. The owning handle a PRIMARY builder
1090 took in begin_building() must then forget that graph rather than destroy
1091 it a second time, and the builder ends up with no graph at all.
1092 """
1093 cdef cydriver.CUgraph c_graph = NULL 2`c~cP Cbnb9bidpbobF lbibjbnd1bqdqbzbL M N O cbH : # x S z rbadQbZ sb3 mb4by *bDbG T U ; ld5 = 6 kbpd4 b c d e f g h i j k l m n o p q A r s t u v B C D E J | ^ ` gbabdb!b8b} _ { hbbbeb+ 7 $ , 8 % 0 V W - 9 ' . ! ( 2 R Q X Y / w 1 ) ? ] Pbfb* a tbAbubvb@ {c|cK I ~ kdIbJbEbKbLbFbMbNbGbObHbwbxbybBb}c[
1094 cdef cydriver.CUresult err = cydriver.CUDA_SUCCESS 2`c~cP Cbnb9bidpbobF lbibjbnd1bqdqbzbL M N O cbH : # x S z rbadQbZ sb3 mb4by *bDbG T U ; ld5 = 6 kbpd4 b c d e f g h i j k l m n o p q A r s t u v B C D E J | ^ ` gbabdb!b8b} _ { hbbbeb+ 7 $ , 8 % 0 V W - 9 ' . ! ( 2 R Q X Y / w 1 ) ? ] Pbfb* a tbAbubvb@ {c|cK I ~ kdIbJbEbKbLbFbMbNbGbObHbwbxbybBb}c[
1095 cdef cydriver.CUstream c_stream
1096 if not gb._h_stream or gb._state != CAPTURING or gb._kind == FORKED: 2`c~cP Cbnb9bidpbobF lbibjbnd1bqdqbzbL M N O cbH : # x S z rbadQbZ sb3 mb4by *bDbG T U ; ld5 = 6 kbpd4 b c d e f g h i j k l m n o p q A r s t u v B C D E J | ^ ` gbabdb!b8b} _ { hbbbeb+ 7 $ , 8 % 0 V W - 9 ' . ! ( 2 R Q X Y / w 1 ) ? ] Pbfb* a tbAbubvb@ {c|cK I ~ kdIbJbEbKbLbFbMbNbGbObHbwbxbybBb}c[
1097 return err 2`c~cP CbnbidpbobF lbibjbnd1bqdqbzbL M N O cbH : # x S z rbadQbZ sb3 mb4by DbG T U ; ld5 = 6 kbpd4 b c d e f g h i j k l m n o p q A r s t u v B C D E J | ^ ` gbabdb!b8b} _ { hbbbeb+ 7 $ , 8 % 0 V W - 9 ' . ! ( 2 R Q X Y / w 1 ) ? ] Pbfb* a tbAbubvb@ {c|cK I ~ kdIbJbEbKbLbFbMbNbGbObHbwbxbybBb}c[
1098 c_stream = as_cu(gb._h_stream) 2P Cb9bpbobF lbibjb1bqbzbL M N O cbH : # x S z rbQbZ sb3 mb4by *bDbG T U ; 5 = 6 kb4 b c d e f g h i j k l m n o p q A r s t u v B C D E J | ^ ` gbabdb!b8b} _ { hbbbeb+ 7 $ , 8 % 0 V W - 9 ' . ! ( 2 R Q X Y / w 1 ) ? ] Pbfb* a tbAbubvb@ K I ~ IbJbEbKbLbFbMbNbGbObHbwbxbybBb[
1099 with nogil: 2P Cb9bpbobF lbibjb1bqbzbL M N O cbH : # x S z rbQbZ sb3 mb4by *bDbG T U ; 5 = 6 kb4 b c d e f g h i j k l m n o p q A r s t u v B C D E J | ^ ` gbabdb!b8b} _ { hbbbeb+ 7 $ , 8 % 0 V W - 9 ' . ! ( 2 R Q X Y / w 1 ) ? ] Pbfb* a tbAbubvb@ K I ~ IbJbEbKbLbFbMbNbGbObHbwbxbybBb[
1100 err = cydriver.cuStreamEndCapture(c_stream, &c_graph) 2P Cb9bpbobF lbibjb1bqbzbL M N O cbH : # x S z rbQbZ sb3 mb4by *bDbG T U ; 5 = 6 kb4 b c d e f g h i j k l m n o p q A r s t u v B C D E J | ^ ` gbabdb!b8b} _ { hbbbeb+ 7 $ , 8 % 0 V W - 9 ' . ! ( 2 R Q X Y / w 1 ) ? ] Pbfb* a tbAbubvb@ K I ~ IbJbEbKbLbFbMbNbGbObHbwbxbybBb[
1101 if c_graph == NULL: 2P Cb9bpbobF lbibjb1bqbzbL M N O cbH : # x S z rbQbZ sb3 mb4by *bDbG T U ; 5 = 6 kb4 b c d e f g h i j k l m n o p q A r s t u v B C D E J | ^ ` gbabdb!b8b} _ { hbbbeb+ 7 $ , 8 % 0 V W - 9 ' . ! ( 2 R Q X Y / w 1 ) ? ] Pbfb* a tbAbubvb@ K I ~ IbJbEbKbLbFbMbNbGbObHbwbxbybBb[
1102 # Only a PRIMARY builder owns a graph hierarchy to retire. A
1103 # CONDITIONAL_BODY graph belongs to the parent's conditional node; the
1104 # driver destroys it as well when the body capture ends invalidated and
1105 # leaves the parent graph referring to it, which cuda.core cannot
1106 # repair (see the end_building docstring). Only this builder's alias
1107 # is dropped here.
1108 if gb._kind == PRIMARY: 2Cb9blbibjbT U
1109 invalidate_root_graph_state(gb._h_graph) 2Cb9blbibjbT U
1110 gb._h_graph.reset() 2Cb9blbibjbT U
1111 gb._state = CAPTURE_INVALIDATED 2Cb9blbibjbT U
1112 else:
1113 gb._state = CAPTURE_ENDED 2P pbobF 1bqbzbL M N O cbH : # x S z rbQbZ sb3 mb4by *bDbG U ; 5 = 6 kb4 b c d e f g h i j k l m n o p q A r s t u v B C D E J | ^ ` gbabdb!b8b} _ { hbbbeb+ 7 $ , 8 % 0 V W - 9 ' . ! ( 2 R Q X Y / w 1 ) ? ] Pbfb* a tbAbubvb@ K I ~ IbJbEbKbLbFbMbNbGbObHbwbxbybBb[
1114 return err 2P Cb9bpbobF lbibjb1bqbzbL M N O cbH : # x S z rbQbZ sb3 mb4by *bDbG T U ; 5 = 6 kb4 b c d e f g h i j k l m n o p q A r s t u v B C D E J | ^ ` gbabdb!b8b} _ { hbbbeb+ 7 $ , 8 % 0 V W - 9 ' . ! ( 2 R Q X Y / w 1 ) ? ] Pbfb* a tbAbubvb@ K I ~ IbJbEbKbLbFbMbNbGbObHbwbxbybBb[
1117cdef inline cydriver.CUresult GB_close(GraphBuilder gb) noexcept:
1118 """Close the builder and return the status of ending its capture.
1120 The builder is closed whatever that status is. close() raises it; join()
1121 attaches it to the error already propagating.
1122 """
1123 cdef cydriver.CUresult err = GB_end_capture(gb) 2`cCbnbF lbibjbqb: # rbZ mby G T U = 6 b c d e f g h i j k l m n o p q r s t u v a tbAbubvb{c|ckdIbJbEbKbLbFbMbNbGbObHbwbxbybBb}c
1124 gb._h_graph.reset() 2`cCbnbF lbibjbqb: # rbZ mby G T U = 6 b c d e f g h i j k l m n o p q r s t u v a tbAbubvb{c|ckdIbJbEbKbLbFbMbNbGbObHbwbxbybBb}c
1125 gb._h_stream.reset() 2`cCbnbF lbibjbqb: # rbZ mby G T U = 6 b c d e f g h i j k l m n o p q r s t u v a tbAbubvb{c|ckdIbJbEbKbLbFbMbNbGbObHbwbxbybBb}c
1126 retry_deferred_cleanup() 2`cCbnbF lbibjbqb: # rbZ mby G T U = 6 b c d e f g h i j k l m n o p q r s t u v a tbAbubvb{c|ckdIbJbEbKbLbFbMbNbGbObHbwbxbybBb}c
1127 gb._state = CLOSED 2`cCbnbF lbibjbqb: # rbZ mby G T U = 6 b c d e f g h i j k l m n o p q r s t u v a tbAbubvb{c|ckdIbJbEbKbLbFbMbNbGbObHbwbxbybBb}c
1128 gb._stream = None 2`cCbnbF lbibjbqb: # rbZ mby G T U = 6 b c d e f g h i j k l m n o p q r s t u v a tbAbubvb{c|ckdIbJbEbKbLbFbMbNbGbObHbwbxbybBb}c
1129 return err 2`cCbnbF lbibjbqb: # rbZ mby G T U = 6 b c d e f g h i j k l m n o p q r s t u v a tbAbubvb{c|ckdIbJbEbKbLbFbMbNbGbObHbwbxbybBb}c
1132cdef inline GraphBuilder GB_init_forked(Stream stream, GraphHandle h_primary_graph):
1133 cdef GraphBuilder gb = GraphBuilder.__new__(GraphBuilder) 1F#ZyGTUbcdefghijklmnopqrstuva
1134 # A FORKED builder captures into the primary's CUgraph. It holds the
1135 # primary's GraphHandle so conditional bodies share its graph hierarchy.
1136 gb._h_graph = h_primary_graph 1F#ZyGTUbcdefghijklmnopqrstuva
1137 gb._h_stream = stream._h_stream 1F#ZyGTUbcdefghijklmnopqrstuva
1138 gb._kind = FORKED 1F#ZyGTUbcdefghijklmnopqrstuva
1139 gb._state = CAPTURING 1F#ZyGTUbcdefghijklmnopqrstuva
1140 gb._stream = stream 1F#ZyGTUbcdefghijklmnopqrstuva
1141 return gb 1F#ZyGTUbcdefghijklmnopqrstuva
1144cdef inline GraphBuilder GB_init_conditional(
1145 Stream stream, cydriver.CUgraph cond_graph,
1146 GraphBuilder parent, cydriver.CUgraphNode owner_node):
1147 cdef GraphBuilder gb = GraphBuilder.__new__(GraphBuilder) 1xSzbcdefghijklmnopqArstuvBCDEwa
1148 gb._h_graph = create_child_graph_handle( 1xSzbcdefghijklmnopqArstuvBCDEwa
1149 cond_graph, parent._h_graph, owner_node)
1150 gb._h_stream = stream._h_stream 1xSzbcdefghijklmnopqArstuvBCDEwa
1151 gb._kind = CONDITIONAL_BODY 1xSzbcdefghijklmnopqArstuvBCDEwa
1152 gb._state = CAPTURE_NOT_STARTED 1xSzbcdefghijklmnopqArstuvBCDEwa
1153 gb._stream = stream 1xSzbcdefghijklmnopqArstuvBCDEwa
1154 return gb 1xSzbcdefghijklmnopqArstuvBCDEwa
1157cdef inline int _get_capture_info(
1158 cydriver.CUstream stream,
1159 cydriver.CUstreamCaptureStatus* status,
1160 cydriver.CUgraph* graph) except?-1 nogil:
1161 """Thin wrapper around ``cuStreamGetCaptureInfo`` that papers over the
1162 CUDA 12 vs 13 signature change.
1164 ``status`` must be non-NULL: the driver rejects ``captureStatus_out=NULL``
1165 with ``CUDA_ERROR_INVALID_VALUE``. ``graph`` may be NULL when the caller
1166 does not need the graph handle.
1167 """
1168 IF CUDA_CORE_BUILD_MAJOR >= 13:
1169 return HANDLE_RETURN(cydriver.cuStreamGetCaptureInfo( 2~cP Cb9bpbobF lbibjb1bqbzbL M N O cbH : # x S z rbQbZ sb3 mb4by *bDbG T U ; 5 = 6 kb4 b c d e f g h i j k l m n o p q A r s t u v B C D E J | ^ ` gbabdb!b8b} _ { hbbbeb+ 7 $ , 8 % 0 V W - 9 ' . ! ( 2 R Q X Y / w 1 ) ? ] Pbfb* a tbAbubvb@ K I ~ IbJbEbKbLbFbMbNbGbObHbwbxbybBb[
1170 stream, status, NULL, graph, NULL, NULL, NULL))
1171 ELSE:
1172 return HANDLE_RETURN(cydriver.cuStreamGetCaptureInfo(
1173 stream, status, NULL, graph, NULL, NULL))
1176cdef inline cydriver.CUgraphNode _capture_tail_node(cydriver.CUstream stream) except *:
1177 """Return the node a freshly-captured single-node operation left as the
1178 stream's sole capture dependency (e.g. the host node added by
1179 ``cuLaunchHostFunc``). The driver advances the stream's dependency set to
1180 the new node, so the next captured op would depend on it.
1181 """
1182 cdef cydriver.CUstreamCaptureStatus status
1183 cdef const cydriver.CUgraphNode* deps = NULL 1LMNOJKI
1184 cdef size_t num_deps = 0 1LMNOJKI
1185 with nogil: 1LMNOJKI
1186 IF CUDA_CORE_BUILD_MAJOR >= 13:
1187 HANDLE_RETURN(cydriver.cuStreamGetCaptureInfo( 1LMNOJKI
1188 stream, &status, NULL, NULL, &deps, NULL, &num_deps))
1189 ELSE:
1190 HANDLE_RETURN(cydriver.cuStreamGetCaptureInfo(
1191 stream, &status, NULL, NULL, &deps, &num_deps))
1192 if num_deps != 1: 1LMNOJKI
1193 raise RuntimeError(
1194 f"expected exactly one capture dependency after a host callback, got {num_deps}")
1195 return <cydriver.CUgraphNode>deps[0] 1LMNOJKI
1198cdef inline tuple GB_cond_with_params(GraphBuilder gb, node_params):
1199 status, _, graph, *deps_info, num_dependencies = handle_return( 2x S z kbb c d e f g h i j k l m n o p q A r s t u v B C D E w a
1200 driver.cuStreamGetCaptureInfo(gb._stream.handle) 2x S z kbb c d e f g h i j k l m n o p q A r s t u v B C D E w a
1201 )
1202 if status != driver.CUstreamCaptureStatus.CU_STREAM_CAPTURE_STATUS_ACTIVE: 2x S z kbb c d e f g h i j k l m n o p q A r s t u v B C D E w a
1203 raise RuntimeError("Cannot add conditional node when not actively capturing") 2kb
1205 new_node = handle_return( 1xSzbcdefghijklmnopqArstuvBCDEwa
1206 driver.cuGraphAddNode(graph, *deps_info, num_dependencies, node_params)) 1xSzbcdefghijklmnopqArstuvBCDEwa
1207 deps_info_update = [[new_node]] + [None] * (len(deps_info) - 1) 1xSzbcdefghijklmnopqArstuvBCDEwa
1209 handle_return( 1xSzbcdefghijklmnopqArstuvBCDEwa
1210 driver.cuStreamUpdateCaptureDependencies( 1xSzbcdefghijklmnopqArstuvBCDEwa
1211 gb._stream.handle, 1xSzbcdefghijklmnopqArstuvBCDEwa
1212 *deps_info_update, # dependencies, edgeData 1xSzbcdefghijklmnopqArstuvBCDEwa
1213 1, # numDependencies
1214 driver.CUstreamUpdateCaptureDependencies_flags.CU_STREAM_SET_CAPTURE_DEPENDENCIES, 1xSzbcdefghijklmnopqArstuvBCDEwa
1215 )
1216 )
1218 return tuple( 1xSzbcdefghijklmnopqArstuvBCDEwa
1219 GB_init_conditional( 1xSzbcdefghijklmnopqArstuvBCDEwa
1220 gb._stream.device.create_stream(), 1xSzbcdefghijklmnopqArstuvBCDEwa
1221 <cydriver.CUgraph><intptr_t>int(node_params.conditional.phGraph_out[i]), 1xSzbcdefghijklmnopqArstuvBCDEwa
1222 gb, 1xSzbcdefghijklmnopqArstuvBCDEwa
1223 <cydriver.CUgraphNode><intptr_t>int(new_node), 1xSzbcdefghijklmnopqArstuvBCDEwa
1224 )
1225 for i in range(node_params.conditional.size) 1xSzbcdefghijklmnopqArstuvBCDEwa
1226 )
1229cdef class Graph:
1230 """An executable graph.
1232 A graph groups a set of CUDA kernels and other CUDA operations together and executes
1233 them with a specified dependency tree. It speeds up the workflow by combining the
1234 driver activities associated with CUDA kernel launches and CUDA API calls.
1236 Graphs must be built using a :obj:`~graph.GraphBuilder` object.
1238 """
1240 def __init__(self):
1241 raise RuntimeError("directly constructing a Graph instance is not supported") 2qd
1243 @staticmethod
1244 cdef Graph _init(GraphExecHandle h_graph_exec):
1245 cdef Graph self = Graph.__new__(Graph) 2P nbF L M N O cbH : x z 3 y G ; 5 = 6 4 b c d e f g h i j k l m n o p q A r s t u v B C D E +b,b-b.b/bSb:bTb#b;b=b?b@bUb[bVb$bWbXb]bYb%bZb^b0b'bwcbdcdxcdded5bscgdyczcAc_b`bhdBctcucCcDcEcFcGcHcvcIcJcKcLcMc{bNcOcPcQcRcScTc| ^ ` gbabdb} _ { hbbbeb+ 7 $ , 8 % 0 V W - 9 ' . ! ( ac|bUcVcWcXcYcZc0c1c2c3c4c5c6cbc:c;c=c?c@c2 [c]c^c/cccdcecfcgchcicjckclcmcncoc_cpc7c8c9c!c#c$c%c'c(c)c*c+c,cRb2b-c.cR Q X Y / qcrc3b6b7bw 1 (b)b) ? ] }b~b* @ K I ~ [
1246 self._h_graph_exec = h_graph_exec 2P nbF L M N O cbH : x z 3 y G ; 5 = 6 4 b c d e f g h i j k l m n o p q A r s t u v B C D E +b,b-b.b/bSb:bTb#b;b=b?b@bUb[bVb$bWbXb]bYb%bZb^b0b'bwcbdcdxcdded5bscgdyczcAc_b`bhdBctcucCcDcEcFcGcHcvcIcJcKcLcMc{bNcOcPcQcRcScTc| ^ ` gbabdb} _ { hbbbeb+ 7 $ , 8 % 0 V W - 9 ' . ! ( ac|bUcVcWcXcYcZc0c1c2c3c4c5c6cbc:c;c=c?c@c2 [c]c^c/cccdcecfcgchcicjckclcmcncoc_cpc7c8c9c!c#c$c%c'c(c)c*c+c,cRb2b-c.cR Q X Y / qcrc3b6b7bw 1 (b)b) ? ] }b~b* @ K I ~ [
1247 return self 2P nbF L M N O cbH : x z 3 y G ; 5 = 6 4 b c d e f g h i j k l m n o p q A r s t u v B C D E +b,b-b.b/bSb:bTb#b;b=b?b@bUb[bVb$bWbXb]bYb%bZb^b0b'bwcbdcdxcdded5bscgdyczcAc_b`bhdBctcucCcDcEcFcGcHcvcIcJcKcLcMc{bNcOcPcQcRcScTc| ^ ` gbabdb} _ { hbbbeb+ 7 $ , 8 % 0 V W - 9 ' . ! ( ac|bUcVcWcXcYcZc0c1c2c3c4c5c6cbc:c;c=c?c@c2 [c]c^c/cccdcecfcgchcicjckclcmcncoc_cpc7c8c9c!c#c$c%c'c(c)c*c+c,cRb2b-c.cR Q X Y / qcrc3b6b7bw 1 (b)b) ? ] }b~b* @ K I ~ [
1249 def close(self) -> None:
1250 """Destroy the graph."""
1251 self._h_graph_exec.reset() 2nb: = 5bsctcucvc|b* K I ~
1252 retry_deferred_cleanup() 2nb: = 5bsctcucvc|b* K I ~
1254 @property
1255 def is_closed(self) -> bool:
1256 """Whether this executable graph has been closed."""
1257 return self._h_graph_exec.get() == NULL 2nb
1259 @property
1260 def handle(self) -> driver.CUgraphExec:
1261 """Return the underlying ``CUgraphExec`` object.
1263 .. caution::
1265 This handle is a Python object. To get the memory address of the underlying C
1266 handle, call ``int()`` on the returned object.
1268 """
1269 return as_py(self._h_graph_exec) 2: /c
1271 @overload
1272 def __getitem__(self, node: KernelNode) -> ExecutableKernelNode: ...
1273 @overload
1274 def __getitem__(self, node: MemsetNode) -> ExecutableMemsetNode: ...
1275 @overload
1276 def __getitem__(self, node: MemcpyNode) -> ExecutableMemcpyNode: ...
1277 @overload
1278 def __getitem__(self, node: ChildGraphNode) -> ExecutableChildGraphNode: ...
1279 @overload
1280 def __getitem__(self, node: EventRecordNode) -> ExecutableEventRecordNode: ...
1281 @overload
1282 def __getitem__(self, node: EventWaitNode) -> ExecutableEventWaitNode: ...
1283 @overload
1284 def __getitem__(self, node: HostCallbackNode) -> ExecutableHostCallbackNode: ...
1286 def __getitem__(self, GraphNode node not None: GraphNode) -> ExecutableGraphNode:
1287 """Return a view for updating *node* in this executable graph.
1289 *node* is a definition node from the graph used to instantiate this
1290 executable. Call ``update()`` on the returned view to replace that
1291 node's parameters for future launches. Kernel, memcpy, and memset
1292 views also support enabling and disabling the node.
1293 """
1294 Graph_check_open(self) 2nbac|bbc:c;c=c?c@c2 [c]c^c/cccdcecfcgchcicjckclcmcncoc_cpcRb2bR Q qcrc3b
1295 GN_check_valid(node) 2ac|bbc:c;c=c?c@c2 [c]c^c/cccdcecfcgchcicjckclcmcncoc_cpcRb2bR Q qcrc3b
1296 return create_executable_node_view( 2ac|bbc:c;c=c?c@c2 [c]c^c/cccdcecfcgchcicjckclcmcncoc_cpcRb2bR Q qcrc3b
1297 self._h_graph_exec, node) 2ac|bbc:c;c=c?c@c2 [c]c^c/cccdcecfcgchcicjckclcmcncoc_cpcRb2bR Q qcrc3b
1299 def update(self, source: "GraphBuilder | GraphDefinition") -> None:
1300 """Update the graph using a new graph definition.
1302 The topology of the provided source must be identical to this graph.
1304 Parameters
1305 ----------
1306 source : :obj:`~graph.GraphBuilder` or :obj:`~graph.GraphDefinition`
1307 The graph definition to update from. A GraphBuilder must have
1308 finished building.
1310 """
1311 Graph_check_open(self) 2nb6 Rb2bR Q X Y 3b6b7bw 1 (b)b) ? ]
1312 cdef GraphHandle h_source
1314 if isinstance(source, GraphBuilder): 26 Rb2bR Q X Y 3b6b7bw 1 (b)b) ? ]
1315 GB_check_open(<GraphBuilder>source) 16w1)?
1316 GB_check_has_graph(<GraphBuilder>source) 1w1)?
1317 if (<GraphBuilder>source)._state != CAPTURE_ENDED: 1w1)?
1318 raise ValueError("Graph has not finished building.") 1?
1319 h_source = (<GraphBuilder>source)._h_graph 1w1)
1320 elif isinstance(source, GraphDefinition): 2Rb2bR Q X Y 3b6b7b(b)b]
1321 GD_check_valid(<GraphDefinition>source) 2Rb2bR Q X Y 3b6b7b(b)b
1322 h_source = (<GraphDefinition>source)._h_graph 2Rb2bR Q X Y 3b6b7b(b)b
1323 else:
1324 raise TypeError( 1]
1325 f"expected GraphBuilder or GraphDefinition, got {type(source).__name__}") 1]
1327 cdef cydriver.CUgraphExecUpdateResultInfo result_info
1328 cdef cydriver.CUresult err = graph_exec_update( 2Rb2bR Q X Y 3b6b7bw 1 (b)b)
1329 self._h_graph_exec, h_source, &result_info)
1330 if err == cydriver.CUresult.CUDA_ERROR_GRAPH_EXEC_UPDATE_FAILURE: 2Rb2bR Q X Y 3b6b7bw 1 (b)b)
1331 reason = driver.CUgraphExecUpdateResult(result_info.result) 2RbQ 6b7b)
1332 msg = f"Graph update failed: {reason.__doc__.strip()} ({reason.name})" 2RbQ 6b7b)
1333 raise CUDAError(msg) 2RbQ 6b7b)
1334 HANDLE_RETURN(err) 22bR X Y 3bw 1 (b)b
1336 def upload(self, Stream stream not None: Stream) -> None:
1337 """Uploads the graph in a stream.
1339 Parameters
1340 ----------
1341 stream : :obj:`~_stream.Stream`
1342 The stream in which to upload the graph
1344 """
1345 Graph_check_open(self) 2nby 5 +b,b-b.b/bSb:bTb;b=b?b@bUb[bVbWbXb]bYbZb^b0b_b`b{b+ 7 $ , 8 % 0 V W - 9 ' . ! ( }b~bI
1346 cdef Stream s = Stream_accept(stream) 2y 5 +b,b-b.b/bSb:bTb;b=b?b@bUb[bVbWbXb]bYbZb^b0b_b`b{b+ 7 $ , 8 % 0 V W - 9 ' . ! ( }b~bI
1347 cdef cydriver.CUgraphExec c_exec = as_cu(self._h_graph_exec) 2y 5 +b,b-b.b/bSb:bTb;b=b?b@bUb[bVbWbXb]bYbZb^b0b_b`b{b+ 7 $ , 8 % 0 V W - 9 ' . ! ( }b~bI
1348 cdef cydriver.CUstream c_stream = as_cu(s._h_stream) 2y 5 +b,b-b.b/bSb:bTb;b=b?b@bUb[bVbWbXb]bYbZb^b0b_b`b{b+ 7 $ , 8 % 0 V W - 9 ' . ! ( }b~bI
1349 with nogil: 2y 5 +b,b-b.b/bSb:bTb;b=b?b@bUb[bVbWbXb]bYbZb^b0b_b`b{b+ 7 $ , 8 % 0 V W - 9 ' . ! ( }b~bI
1350 HANDLE_RETURN(cydriver.cuGraphUpload(c_exec, c_stream)) 2y 5 +b,b-b.b/bSb:bTb;b=b?b@bUb[bVbWbXb]bYbZb^b0b_b`b{b+ 7 $ , 8 % 0 V W - 9 ' . ! ( }b~bI
1352 def launch(self, Stream stream not None: Stream) -> None:
1353 """Launches the graph in a stream.
1355 Parameters
1356 ----------
1357 stream : :obj:`~_stream.Stream`
1358 The stream in which to launch the graph.
1360 """
1361 Graph_check_open(self) 2P nbL M N O H x z 3 y ; 5 4 b c d e f g h i j k l m n o p q A r s t u v B C D E +b,b-b.b/bSb:bTb#b;b=b?b@bUb[bVb$bWbXb]bYb%bZb^b0b'b5bscyczcAc_b`bBctcucCcDcEcFcGcHcvcIcJcKcLcMc{bNcOcPcQcRcScTc+ 7 $ , 8 % 0 V W - 9 ' . ! ( ac|bUcVcWcXcYcZc0c1c2c3c4c5c6cbc2 ccdcecfcgchcicjckclcmcncocpc7c8c9c!c#c$c%c'c(c)c*c+c,cRb2b-c.cR Q X Y / qcrc3b6b7bw 1 (b)b}b~b@ K I [
1362 cdef Stream s = Stream_accept(stream) 2P L M N O H x z 3 y ; 5 4 b c d e f g h i j k l m n o p q A r s t u v B C D E +b,b-b.b/bSb:bTb#b;b=b?b@bUb[bVb$bWbXb]bYb%bZb^b0b'b5bscyczcAc_b`bBctcucCcDcEcFcGcHcvcIcJcKcLcMc{bNcOcPcQcRcScTc+ 7 $ , 8 % 0 V W - 9 ' . ! ( ac|bUcVcWcXcYcZc0c1c2c3c4c5c6cbc2 ccdcecfcgchcicjckclcmcncocpc7c8c9c!c#c$c%c'c(c)c*c+c,cRb2b-c.cR Q X Y / qcrc3b6b7bw 1 (b)b}b~b@ K I [
1363 cdef cydriver.CUgraphExec c_exec = as_cu(self._h_graph_exec) 2P L M N O H x z 3 y ; 5 4 b c d e f g h i j k l m n o p q A r s t u v B C D E +b,b-b.b/bSb:bTb#b;b=b?b@bUb[bVb$bWbXb]bYb%bZb^b0b'b5bscyczcAc_b`bBctcucCcDcEcFcGcHcvcIcJcKcLcMc{bNcOcPcQcRcScTc+ 7 $ , 8 % 0 V W - 9 ' . ! ( ac|bUcVcWcXcYcZc0c1c2c3c4c5c6cbc2 ccdcecfcgchcicjckclcmcncocpc7c8c9c!c#c$c%c'c(c)c*c+c,cRb2b-c.cR Q X Y / qcrc3b6b7bw 1 (b)b}b~b@ K I [
1364 cdef cydriver.CUstream c_stream = as_cu(s._h_stream) 2P L M N O H x z 3 y ; 5 4 b c d e f g h i j k l m n o p q A r s t u v B C D E +b,b-b.b/bSb:bTb#b;b=b?b@bUb[bVb$bWbXb]bYb%bZb^b0b'b5bscyczcAc_b`bBctcucCcDcEcFcGcHcvcIcJcKcLcMc{bNcOcPcQcRcScTc+ 7 $ , 8 % 0 V W - 9 ' . ! ( ac|bUcVcWcXcYcZc0c1c2c3c4c5c6cbc2 ccdcecfcgchcicjckclcmcncocpc7c8c9c!c#c$c%c'c(c)c*c+c,cRb2b-c.cR Q X Y / qcrc3b6b7bw 1 (b)b}b~b@ K I [
1365 with nogil: 2P L M N O H x z 3 y ; 5 4 b c d e f g h i j k l m n o p q A r s t u v B C D E +b,b-b.b/bSb:bTb#b;b=b?b@bUb[bVb$bWbXb]bYb%bZb^b0b'b5bscyczcAc_b`bBctcucCcDcEcFcGcHcvcIcJcKcLcMc{bNcOcPcQcRcScTc+ 7 $ , 8 % 0 V W - 9 ' . ! ( ac|bUcVcWcXcYcZc0c1c2c3c4c5c6cbc2 ccdcecfcgchcicjckclcmcncocpc7c8c9c!c#c$c%c'c(c)c*c+c,cRb2b-c.cR Q X Y / qcrc3b6b7bw 1 (b)b}b~b@ K I [
1366 HANDLE_RETURN(cydriver.cuGraphLaunch(c_exec, c_stream)) 2P L M N O H x z 3 y ; 5 4 b c d e f g h i j k l m n o p q A r s t u v B C D E +b,b-b.b/bSb:bTb#b;b=b?b@bUb[bVb$bWbXb]bYb%bZb^b0b'b5bscyczcAc_b`bBctcucCcDcEcFcGcHcvcIcJcKcLcMc{bNcOcPcQcRcScTc+ 7 $ , 8 % 0 V W - 9 ' . ! ( ac|bUcVcWcXcYcZc0c1c2c3c4c5c6cbc2 ccdcecfcgchcicjckclcmcncocpc7c8c9c!c#c$c%c'c(c)c*c+c,cRb2b-c.cR Q X Y / qcrc3b6b7bw 1 (b)b}b~b@ K I [
1369cdef inline int Graph_check_open(Graph self) except -1:
1370 if not self._h_graph_exec: 2P nbL M N O H x z 3 y ; 5 6 4 b c d e f g h i j k l m n o p q A r s t u v B C D E +b,b-b.b/bSb:bTb#b;b=b?b@bUb[bVb$bWbXb]bYb%bZb^b0b'b5bscyczcAc_b`bBctcucCcDcEcFcGcHcvcIcJcKcLcMc{bNcOcPcQcRcScTc+ 7 $ , 8 % 0 V W - 9 ' . ! ( ac|bUcVcWcXcYcZc0c1c2c3c4c5c6cbc:c;c=c?c@c2 [c]c^c/cccdcecfcgchcicjckclcmcncoc_cpc7c8c9c!c#c$c%c'c(c)c*c+c,cRb2b-c.cR Q X Y / qcrc3b6b7bw 1 (b)b) ? ] }b~b@ K I [
1371 raise RuntimeError("Graph has been closed") 2nb
1372 return 0 2P L M N O H x z 3 y ; 5 6 4 b c d e f g h i j k l m n o p q A r s t u v B C D E +b,b-b.b/bSb:bTb#b;b=b?b@bUb[bVb$bWbXb]bYb%bZb^b0b'b5bscyczcAc_b`bBctcucCcDcEcFcGcHcvcIcJcKcLcMc{bNcOcPcQcRcScTc+ 7 $ , 8 % 0 V W - 9 ' . ! ( ac|bUcVcWcXcYcZc0c1c2c3c4c5c6cbc:c;c=c?c@c2 [c]c^c/cccdcecfcgchcicjckclcmcncoc_cpc7c8c9c!c#c$c%c'c(c)c*c+c,cRb2b-c.cR Q X Y / qcrc3b6b7bw 1 (b)b) ? ] }b~b@ K I [