View page source

warp.tile_stack_push#

warp.tile_stack_push(
s: Any,
value: Any,
has_value: bool,
) → int#

  • Kernel

Push a value onto a tile stack (cooperative).

All threads in the block must call this function. Only threads with has_value=True write to the stack.

Parameters:

  • s – The tile stack

  • value – The value to push

  • has_value – Whether this thread has a value to push

Returns:

The slot index where the value was written, or -1 if has_value is False or the stack overflowed.

Example

CAP = wp.constant(8)

@wp.kernel
def push_kernel(out_idx: wp.array[int]):
    _i, j = wp.tid()
    s = wp.tile_stack(capacity=CAP, dtype=int)
    idx = wp.tile_stack_push(s, j * 10, j < 4)
    out_idx[j] = idx

out_idx = wp.full(8, -1, dtype=int)
wp.launch_tiled(push_kernel, dim=[1], inputs=[out_idx], block_dim=8)

idxs = out_idx.numpy()
print(sorted(idxs[idxs >= 0].tolist()))
print(sum(idxs == -1))

[0, 1, 2, 3]
4