CUDAGraph

class tensorplay.cuda.graphs.CUDAGraph[source]

begin_capture_to_if_node(scalar_pred)[source]

Inside an open capture, gate the following work on an if node.

scalar_pred must be a single-element CUDA Bool tensor; at replay time the driver samples it and runs the body captured between this call and end_capture_to_conditional_node() only when true.

begin_capture_to_while_node(scalar_pred)[source]

Like begin_capture_to_if_node(), but the body loops while the predicate stays true (driver-level while node).

capture_begin(pool=None, capture_error_mode='global', stream=None)[source]

Begin capture.

Parameters:

• pool – None captures into a fresh private pool; an id from graph_pool_handle(), another graph’s pool_id, or another CUDAGraph shares that pool instead.

• capture_error_mode (str) – "global" fails the capture if any unsafe CUDA call happens anywhere in the process; "thread_local" only watches this thread; "relaxed" does not guard against unsafe calls.

• stream – custom capture stream (a tensorplay.cuda.Stream). Defaults to the runtime’s dedicated per-device side stream; the legacy default stream cannot participate in capture.

capture_end()[source]

End capture and compile the executable (paid here, not on first replay).

debug_dump(path)[source]

Write a DOT rendering of the captured graph to path.

Call enable_debug_mode() before capturing for a dump that includes full node attributes.

end_capture_to_conditional_node()[source]

Close the open conditional body; subsequent capture returns to the parent stream.

instantiate()[source]

No-op once instantiated; kept for late callers.

property pool_id

Allocator pool id this graph captured against.

replay(stream=None)[source]

Run the graph: launch the cached executable on the current stream.

Parameters:

stream (Stream, optional) – launch on this explicit stream instead of querying the current one - shaves a TLS lookup off hot loops pinned to a single stream.

reset()[source]

Destroy the executable and release the pool reference.

All tensors allocated during the capture must be released first.

set_conditional_handle_for_current_node(scalar_pred)[source]

Refresh the predicate consumed by the innermost open conditional node (used for nested conditionals).

stage_and_launch(static_inputs, inputs)[source]

Stage every input onto its static buffer and replay in one call.

Parameters:

• static_inputs – buffers captured by the graph (kept alive by the caller).

• inputs – fresh tensors whose contents overwrite the matching static buffer this iteration. Contiguous same-dtype/ same-device pairs take a raw async device-to-device copy; anything else falls back to full copy semantics.

This is the low-overhead bulk entry used by tensorplay._stax.cudagraphs: one Python-to-native crossing for the whole replay instead of one dispatcher round trip per input.