@triton.jit
def union_kernel(parent, rank, u, v, BLOCK_SIZE: tl.constexpr):
    root_u = tl.zeros((BLOCK_SIZE,), dtype=tl.int32)
    root_v = tl.zeros((BLOCK_SIZE,), dtype=tl.int32)
    find_kernel(parent, u, root_u, BLOCK_SIZE=BLOCK_SIZE)
    find_kernel(parent, v, root_v, BLOCK_SIZE=BLOCK_SIZE)
    tl.syncwarp()
    root_u = root_u[0]
    root_v = root_v[0]
    if root_u != root_v:
        ru_rank = tl.load(rank + root_u)
        rv_rank = tl.load(rank + root_v)
        if ru_rank > rv_rank:
            tl.store(parent, root_v, root_u)
        elif ru_rank < rv_rank:
            tl.store(parent, root_u, root_v)
        else:
            tl.store(parent, root_v, root_u)
            tl.atomic_add(rank, root_u, 1)
