warp.tile_argmax

warp.tile_argmax(

a: Tile[Scalar, tuple[int, ...]],

) → Tile[Int, tuple[1]]

Kernel

Cooperatively compute the index of the maximum element in the tile using all threads in the block.

param a:

The tile to compute the argmax from

returns:

A single-element tile holding the index of the maximum value

Example:

@wp.kernel
def compute():

    t = wp.tile_arange(64, 128)
    s = wp.tile_argmax(t)

    print(s)

wp.launch_tiled(compute, dim=[1], inputs=[], block_dim=64)

Prints:

[63] = tile(shape=(1), storage=register)