Function _odd_even_sort

python/tvm/topi/gpu/sort.py:85–180

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

def _odd_even_sort(
    size,
    axis_mul_before,
    axis_mul_after,
    is_ascend,
    keys,
    keys_swap,
    values=None,
    values_swap=None,
):
    nthread_tx = block_size // 2
    nthread_bx = ceil_div(size, block_size)
    nthread_by = axis_mul_before * axis_mul_after

    tx, bx, by, ntx, nbx, nby = _get_threads(nthread_tx, nthread_bx, nthread_by)
    with T.frame_scope(
        [
            T.attr(tvm.tirx.const(0), "hand_threaded", 0),
            T.attr(tx, "thread_extent", ntx),
            T.attr(bx, "thread_extent", nbx),
            T.attr(by, "thread_extent", nby),
        ]
    ):
        by_val = by % axis_mul_before
        bz = by // axis_mul_before
        tid = 2 * tx
        start = bx * block_size

        # Buffer declarations (DeclBuffer generates both Allocate + DeclBuffer nodes)
        tmp_keys_swap = T.decl_buffer([block_size], keys_swap.dtype, scope="shared")
        temp_keys = T.decl_buffer([1], keys_swap.dtype, scope="local")
        temp_cond1 = T.decl_buffer([1], keys_swap.dtype, scope="local")
        temp_cond2 = T.decl_buffer([1], keys_swap.dtype, scope="local")
        if values_swap is not None:
            tmp_values_swap = T.decl_buffer([block_size], values_swap.dtype, scope="shared")
            temp_values = T.decl_buffer([1], values_swap.dtype, scope="local")

        # Copy data to scratch space
        base_idx = by_val * size * axis_mul_after + bz
        with T.serial(0, 2) as n:
            with T.If((tid + n + start) < size):
                with T.Then():
                    T.buffer_store(
                        tmp_keys_swap,
                        keys[base_idx + (tid + n + start) * axis_mul_after],
                        [tid + n],
                    )
                    if values_swap is not None:
                        T.buffer_store(
                            tmp_values_swap,
                            values[base_idx + (tid + n + start) * axis_mul_after],
                            [tid + n],
                        )

        T.evaluate(tvm.tirx.Call(None, "tirx.tvm_storage_sync", tvm.runtime.convert(["shared"])))

        idxm = tvm.tirx.indexmod
        # OddEvenTransposeSort