warp.utils.segmented_sort_pairs#

warp.utils.segmented_sort_pairs(
keys,
values,
count,
segment_start_indices,
segment_end_indices=None,
)[source]#

Sort key-value pairs in place within each segment in ascending key order.

The sort is stable within each segment. Pairs with identical keys retain their original relative order. Elements in keys[:count] and values[:count] outside the segment ranges are not modified.

Segments are half-open ranges defined by segment_start_indices and optionally segment_end_indices. Every range must satisfy 0 <= start <= end <= count.

During direct CPU execution, Warp raises a ValueError for invalid segment bounds. If invalid bounds are encountered during CPU graph replay, warp.capture_launch() raises a RuntimeError. For execution on a CUDA device, Warp does not currently report invalid segment bounds, so callers must validate segment ranges before calling this function.

All provided arrays must be contiguous and reside on the same device. The segment-index arrays must be one-dimensional.

The keys and values arrays must each contain at least 2 * count elements. The keys[count:2 * count] and values[count:2 * count] regions are scratch storage and may be overwritten.

Segment ranges must not overlap one another, and neither segment_start_indices nor segment_end_indices may overlap the storage used by keys or values. Warp does not currently detect either type of overlap. Violating either restriction results in undefined behavior.

Parameters:

  • keys (warp.array[warp.int32] | warp.array[warp.float32]) – Array of keys to sort. Its dtype must be int32 or float32.

  • values (warp.array[warp.int32]) – Array of values to reorder with their corresponding keys. Its dtype must be int32.

  • count (int) – Number of elements available to the segment ranges at the start of keys and values. Must satisfy 0 <= count <= 2**31 - 1.

  • segment_start_indices (warp.array[warp.int32]) – Start index of each segment. When segment_end_indices is None, adjacent entries define each segment, so an array of length N defines N - 1 segments.

  • segment_end_indices (warp.array[warp.int32] | None) – End index of each segment. When provided, it must have the same length as segment_start_indices.

Raises:

  • TypeError – If count is not an integer.

  • ValueError – If count is outside the signed 32-bit non-negative integer range or a segment range is invalid during direct CPU execution.

  • RuntimeError – If the arrays reside on different devices, keys or values has insufficient storage, an array is not contiguous, a segment-index array is not one-dimensional, the segment-index arrays have different lengths, or a dtype is unsupported.

Return type:

None

Example

>>> keys = wp.array([3, 1, 4, 2, 0, 0, 0, 0], dtype=wp.int32)
>>> values = wp.array([30, 10, 40, 20, 0, 0, 0, 0], dtype=wp.int32)
>>> offsets = wp.array([0, 2, 4], dtype=wp.int32)
>>> wp.utils.segmented_sort_pairs(keys, values, 4, offsets)
>>> keys.numpy()[:4].tolist()
[1, 3, 2, 4]
>>> values.numpy()[:4].tolist()
[10, 30, 20, 40]
On this page