Function _bitonic_merge

src/fused_mm_sampling/tl_argsort.py:55–62

(x, ids, stage: core.constexpr, order: core.constexpr, n_dims: core.constexpr)

@triton.jit
def _bitonic_merge(x, ids, stage: core.constexpr, order: core.constexpr, n_dims: core.constexpr):
    if order == 2:
        flip = _indicator(n_dims, stage)
    else:
        flip = order
    for i in core.static_range(stage):
        x, ids = _compare_and_swap(x, ids, flip, stage - 1 - i, n_dims)
    return x, ids


@triton.jit