warp.ScopedCapture#

class warp.ScopedCapture( device=None, stream=None, force_module_load=None, external=False, )[source]#

Context manager to capture a sequence of operations into a CUDA graph.

CUDA graphs allow a sequence of GPU operations to be captured and replayed with reduced launch overhead. The captured graph is available as the graph attribute after exiting the context.

Parameters:

device (DeviceLike) – Device on which to capture operations.
stream – Stream on which to capture operations.
force_module_load – If True, force all modules to load before capture begins.
external – If True, indicates an external graph capture is already active.

graph#: The captured CUDA graph, available after context exit.

active#: Whether capture is currently in progress.

Example

with wp.ScopedCapture() as capture:
    wp.launch(my_kernel, dim=n, inputs=[arr])

# Replay the captured graph
wp.capture_launch(capture.graph)