warp.capture_if#

warp.capture_if(
condition,
on_true=None,
on_false=None,
stream=None,
**kwargs,
)[source]#

Create a dynamic branch based on a condition.

The condition value is retrieved from the first element of the condition array.

This function is particularly useful with CUDA graphs, but can be used without graph capture as well. CUDA 12.4+ is required for CUDA graph conditional nodes. During APIC recording (all CPU captures and CUDA captures with apic=True), branch bodies must be callbacks; passing Graph objects is not yet supported. CUDA capture with apic=False supports either form.

Parameters:

  • condition (warp.array[warp.int32]) – Warp array holding the condition value.

  • on_true (Callable | Graph | None) – A callback function or Graph to execute if the condition is True.

  • on_false (Callable | Graph | None) – A callback function or Graph to execute if the condition is False.

  • stream (Stream) – The CUDA stream where the condition was written. If None, use the current stream on the device where condition resides.

Any additional keyword arguments are forwarded to the callback functions.

On this page