warp.tile_sort#

warp.tile_sort( keys: Tile[Any, tuple[int, ...]], values: Tile[Any, tuple[int, ...]], ) → None#

Kernel: true
Python: false
Differentiable: false

Sort two tiles in place, in ascending order of keys.

This is a cooperative operation: the threads of the block divide the work between them, so every thread must reach the call.

Both tiles are permuted together, so values acts as a payload that follows its key. The sort is not stable: the relative order of equal keys is unspecified and differs between CPU and GPU. Multi-dimensional tiles are sorted as their flattened row-major sequence. Both tiles are placed in shared memory whatever storage they were created with, so they count against the block’s shared-memory budget (see Shared Memory Budget).

On CUDA devices, prefer a block_dim that is a multiple of 32 for better sorting performance. Other block dimensions are supported but may be slower.

Parameters:

keys – Tile of sort keys, permuted in place. Supported key types: warp.float32, warp.int32, warp.uint32, warp.int64, warp.uint64.
values – Payload tile, permuted in place. Must have the same shape as keys; any element type is accepted.

Example

@wp.kernel
def sort_pairs(k: wp.array[float], v: wp.array[int]):
    keys = wp.tile_load(k, shape=4)
    values = wp.tile_load(v, shape=4)
    wp.tile_sort(keys, values)
    wp.tile_store(k, keys)
    wp.tile_store(v, values)

keys = wp.array([5.0, 3.0, 8.0, 1.0], dtype=float)
values = wp.array([0, 1, 2, 3], dtype=int)

wp.launch_tiled(sort_pairs, dim=1, outputs=[keys, values], block_dim=32)

print(keys.numpy())
print(values.numpy())

[1. 3. 5. 8.]
[3 1 0 2]