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_: UnsafePointer[Scalar[dtype], MutUntrackedOrigin], output_keys_: UnsafePointer[Scalar[dtype], MutUntrackedOrigin], input_key_ids_: UnsafePointer[Scalar[out_idx_type], MutUntrackedOrigin], output_key_ids_: UnsafePointer[Scalar[out_idx_type], MutUntrackedOrigin], num_keys: Int, skip_sort: UnsafePointer[Scalar[DType.bool], MutUntrackedOrigin])

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_ (UnsafePointer[Scalar[dtype], MutUntrackedOrigin]): Input tensor values to sort.
output_keys_ (UnsafePointer[Scalar[dtype], MutUntrackedOrigin]): Output tensor values sorted in (default) descending order.
input_key_ids_ (UnsafePointer[Scalar[out_idx_type], MutUntrackedOrigin]): Input tensor indices.
output_key_ids_ (UnsafePointer[Scalar[out_idx_type], MutUntrackedOrigin]): Output tensor indices sorted in (default) descending order.
num_keys (Int): Number of keys to sort per batch.
skip_sort (UnsafePointer[Scalar[DType.bool], MutUntrackedOrigin]): Whether sorting is skipped for this batch.