warp.capture_begin#

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

Begin capture of a graph.

Records supported Warp operations subsequently issued on the selected device. On CUDA devices, the CUDA driver captures operations into a CUDA graph. On CPU devices, operations are recorded into an API Capture (APIC) operation stream, which capture_launch() replays in C++.

If apic=True, APIC recording is also performed alongside CUDA graph capture, and the result can be serialized to a .wrp file via capture_save(). The flag has no effect on CPU (recording is always on there) beyond gating whether capture_save() is allowed.

Parameters:

  • device (Device | str | None) – The device to capture on (CUDA or CPU).

  • stream (Stream | None) – The CUDA stream to capture on (CUDA only).

  • force_module_load (bool | None) – Whether to force loading of all kernels before capture. In general it is better to use load_module() to selectively load kernels. When running with CUDA drivers that support CUDA 12.3 or newer, this option is not recommended to be set to True because kernels can be loaded during graph capture on more recent drivers. If this argument is None, then the behavior inherits from warp.config.enable_graph_capture_module_load_by_default if the driver is older than CUDA 12.3.

  • external (bool) – Whether the capture was already started externally (CUDA only). The capture_mode argument should specify the mode that was used to initiate the external capture.

  • apic (bool) – Whether to allow capture_save() on the captured graph. On CUDA this also enables APIC operation-stream recording during the capture; on CPU, recording happens regardless because it is the only replay mechanism.

  • capture_mode (CaptureMode) – The CaptureMode (i.e. cudaStreamCaptureMode) used when Warp opens the capture. Defaults to CaptureMode.THREAD_LOCAL. Use CaptureMode.RELAXED when composing with libraries that may call capture-unsafe CUDA runtime APIs during the capture (e.g. lazy context initialization). Ignored on CPU devices.

On this page