Function _sort_common

python/tvm/topi/gpu/sort.py:183–578

Either sort only values or sort values by keys.

(
    size,
    axis_mul_before,
    axis_mul_after,
    is_ascend,
    keys,
    keys_swap,
    values=None,
    values_swap=None,
)

def _sort_common(
    size,
    axis_mul_before,
    axis_mul_after,
    is_ascend,
    keys,
    keys_swap,
    values=None,
    values_swap=None,
):
    """Either sort only values or sort values by keys."""

    ## This function performs a multi-level mergesort
    ## For blocks of length <= block_size, it does odd-even transpose sort
    ##    in GPU shared memory
    ## For intermediate block sizes (>block_size, < max_threads * thread_work)
    ##    it uses the mergpath algorthim https://arxiv.org/abs/1406.2628
    ##    to merge blocks in parallel
    ## At some point, the size of the blocks to be merged is too big for max_threads
    ##    and we switch to using a dual-level mergepath where the outer mergepath
    ##    finds the start/end locations of the inner mergepath so that we can split
    ##    the merge into more blocks

    target = tvm.target.Target.current(allow_none=False)
    max_threads = int(target.attrs["max_num_threads"])
    is_webgpu = "webgpu" in str(target)
    target_dtype = "int32" if is_webgpu else "int64"
    nthread_by = axis_mul_before * axis_mul_after
    nthread_tx = max_threads
    nthread_bx = ceil_div(size, nthread_tx)

    def compare(a, b):
        """Compare a and b in proper ascending or descending order"""
        if is_ascend:
            out = a <= b
        else:
            out = b <= a
        return out

    # Sort the lower levels of the merge using odd-even sort, it's fast for small inputs
    lower_lim = ceil_log2(block_size)

    _odd_even_sort(
        size,
        axis_mul_before * axis_mul_after,
        1,
        is_ascend,
        keys,
        keys_swap,
        values,
        values_swap,
    )

    upper_lim = ceil_log2(size)

    def get_merge_begin(source, base_idx, aCount, bCount, aStart, bStart, diag, first, last):
        max_val = tvm.te.max(0, diag - bCount)
        min_val = tvm.te.min(diag, aCount)