View page source

warp.capture_while#

warp.capture_while(condition, while_body, stream=None, **kwargs)[source]#

Create a dynamic loop based on a condition.

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

The while_body callback is responsible for updating the condition value so the loop can terminate.

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. CPU APIC capture records this operation directly, but the loop body must be a callback; passing a Graph object under CPU APIC capture is not yet supported.

Parameters:

  • condition (warp.array(dtype=wp.int32, ndim=1)) – Warp array holding the condition value.

  • while_body (Callable | Graph) – A callback function or Graph to execute while the loop condition is True.

  • stream (Stream | None) – 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 function.