Asynchronous Copies

This guide shows how to copy data asynchronously and wait before using it. A kernel can issue a transfer, perform independent work, then wait for completion. The copy instruction and its completion calls are separate parts of that sequence.

A Global-to-Shared Copy

cp.async copies a short contiguous chunk directly from global to shared memory. On an Ampere-or-later GPU, we can use it for the row exchange from Data Movement. Each thread copies one four-float row, then reads the next thread’s row:

from tvm_tirx.script import tirx as T

@T.function
def async_read_next_row(
    A: T.Tensor((64, 4), "float32", align=16),
    B: T.Tensor((64, 4), "float32", align=16),
):
    T.device_entry(launch=T.cuda.LaunchConfig(grid=1, block=64))
    tid = T.cuda.thread_idx("x")
    shared = T.alloc_shared((64, 4), "float32", align=16)
    T.cuda.tile.apply(
        T.ptx["cp.async.cg.shared.global"],
        shared.sub[tid, :], A.sub[tid, :], chunk_bytes=16,
    )
    T.ptx.cp.async_.commit_group()
    T.ptx.cp.async_.wait_group(0)
    T.cuda.cta_sync()
    for j in range(4):
        B[tid, j] = shared[(tid + 1) % 64, j]

The source and destination views each contain four float32 elements, which occupy the requested 16 bytes. Each thread issues one copy. The following calls complete the handoff:

1. commit_group() collects that thread’s previously uncommitted copies into a copy group.

2. wait_group(0) waits for all of its committed groups to complete.

3. cta_sync() lets the threads read data copied by their neighbors.

Independent computation can run between the commit and wait. A CTA barrier alone does not wait for outstanding cp.async transfers.

We can check that the output has the same row shift as the synchronous version:

import torch
import tvm_tirx

a = torch.arange(256, dtype=torch.float32, device="cuda").reshape(64, 4)
b = torch.empty_like(a)
tvm_tirx.compile(async_read_next_row, target="cuda")(a, b)
torch.testing.assert_close(b, torch.roll(a, shifts=-1, dims=0))

A false pred skips the copy and leaves the destination unchanged. src_size instead copies a byte prefix and zero-fills the rest, while ignore_src=True zero-fills the whole chunk. The declared source region must still satisfy the bounds checks. The cp.async reference lists supported widths and optional operands.

Copies between CTAs

A thread-block cluster can also transfer shared memory between its CTAs. The cp.async.bulk.shared::cluster.shared::cta form uses a transaction barrier in the destination CTA. Both CTAs must have matching live shared allocations and initialize their barriers before the sender issues a copy.

The sender passes its local source address and the destination’s remote shared and barrier addresses to the copy instruction. For a 256-byte chunk, the receiver registers 256 expected transaction bytes and waits for its barrier phase before reading the destination. Both CTAs synchronize before either exits or reuses the participating storage.

This path uses the destination’s transaction barrier rather than ordinary cp.async copy groups. The complete cluster example uses raw PTX to show address mapping, proxy fences, expected-byte accounting, and the receiver’s wait. TMA’s descriptor-based global/shared transfers are covered separately in TMA Transfers.