For the complete documentation index, see llms.txt. Markdown versions of all pages are available by appending .md to any URL (e.g. /max/get-started.md).

Mojo function

radix_sort_pairs_kernel

def radix_sort_pairs_kernel[dtype: DType, out_idx_type: DType, current_bit: Int, ascending: Bool = False, BLOCK_SIZE: Int = Int(256), NUM_BITS_PER_PASS: Int = Int(4)](input_keys_: Pointer[Scalar[dtype], MutUntrackedOrigin, _safe=False], output_keys_: Pointer[Scalar[dtype], MutUntrackedOrigin, _safe=False], input_key_ids_: Pointer[Scalar[out_idx_type], MutUntrackedOrigin, _safe=False], output_key_ids_: Pointer[Scalar[out_idx_type], MutUntrackedOrigin, _safe=False], num_keys: Int, skip_sort: Pointer[Scalar[DType.bool], MutUntrackedOrigin, _safe=False])

Radix pair sort kernel for (default) descending order.

Implementation based on: AMD. Introduction to GPU Radix Sort. GPUOpen, 2017. Available at: https://gpuopen.com/download/publications/Introduction_to_GPU_Radix_Sort.pdf.

Parameters:

dtype (DType): DType - Data type.
out_idx_type (DType): DType - Output index type.
current_bit (Int): Int - Current bit to start sorting NUM_BITS_PER_PASS bits at.
ascending (Bool): Bool - Whether to sort in ascending order.
BLOCK_SIZE (Int): Int - Block size.
NUM_BITS_PER_PASS (Int): Int - Number of bits per pass.

Args:

input_keys_ (Pointer[Scalar[dtype], MutUntrackedOrigin, _safe=False]): Input tensor values to sort.
output_keys_ (Pointer[Scalar[dtype], MutUntrackedOrigin, _safe=False]): Output tensor values sorted in (default) descending order.
input_key_ids_ (Pointer[Scalar[out_idx_type], MutUntrackedOrigin, _safe=False]): Input tensor indices.
output_key_ids_ (Pointer[Scalar[out_idx_type], MutUntrackedOrigin, _safe=False]): Output tensor indices sorted in (default) descending order.
num_keys (Int): Number of keys to sort per batch.
skip_sort (Pointer[Scalar[DType.bool], MutUntrackedOrigin, _safe=False]): Whether sorting is skipped for this batch.