Mojo function

radix_sort_pairs_kernel

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

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