warp.ScopedCapture#

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

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

Captures kernel launches, memory copies, and memsets for later replay with reduced launch overhead. Works on both CPU and CUDA devices. The captured graph is available as the graph attribute after exiting the context.

Parameters:

device (DeviceLike) – Device on which to capture operations (CPU or CUDA).
stream – Stream on which to capture operations (CUDA only).
force_module_load – If True, force all modules to load before capture begins.
external – If True, indicates an external graph capture is already active. The capture_mode argument should specify the mode that was used to initiate the external capture.
apic (bool) – If True, enable APIC recording for serialization via capture_save(). On CPU, recording always occurs regardless of this flag (needed for CPU graph replay). Default is False.
capture_mode (CaptureMode) – The CaptureMode to use when opening the capture. Defaults to CaptureMode.THREAD_LOCAL. See capture_begin() for details.

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

active#: Whether capture is currently in progress.

Example

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

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