copy_async → ldgsts#

The ldgsts variant lowers copy_async for a global → shared transfer to the PTX cp.async (LDGSTS) instruction: each thread issues an asynchronous vectorized copy that the hardware completes in the background, so the warp can keep computing while the load is in flight. It reuses the exact [outer, threads, vec] partition of the synchronous copy → vec_auto global/shared path vec_auto path; the differences are all in what is emitted and when it completes. Source: python/tvm/backend/cuda/tile_primitive/copy_async/ldgsts.py.

What it accepts#

_LDGSTS_PAIRS = [("global", "shared*")]      # cp.async is unidirectional
_LDGSTS_VEC_BITS = (128, 64, 32)             # cp_size ∈ {16, 8, 4} bytes

def _is_ldgsts(op_call, sctx):
    if not sctx.is_target("cuda"):
        return False, "non-cuda target"
    if sctx.scope_kind not in ("thread", "warp", "warpgroup", "cta"):
        return False, f"unsupported exec_scope {sctx.scope_kind}"
    for check in (
        lambda: _all_threads_active(sctx),
        lambda: _is_valid_copy(op_call, sctx),
        lambda: _scope_allowed(op_call, sctx, allowed_pairs=_LDGSTS_PAIRS),
        lambda: _divides_thread_cnt_ldgsts(op_call, sctx),
    ):
        ok, msg = check()
        if not ok:
            return False, msg
    return True, None

Property	Requirement
target / scope	`cuda`; `thread` / `warp` / `warpgroup` / `cta`, all active
direction	global → shared only (`_LDGSTS_PAIRS`) — `cp.async` is unidirectional
dtype / shape	`_is_valid_copy`; region element count is divisible by the thread count
vector width	`cp.async` only accepts cp_size ∈ {4, 8, 16} bytes, so the candidate set is restricted to `_LDGSTS_VEC_BITS` = `{128, 64, 32}` bits
priority	`20` — selected for `copy_async` global→shared (vs the bulk/TMA variants)

Configuration#

ldgsts reads four optional config keys. prefetch_size is -1 by default (no qualifier), or 64, 128, or 256 for the corresponding PTX .L2::<N>B qualifier. predicate defaults to -1 (unpredicated); with fill_mode="zero", a false predicate uses the src-size form to zero-fill the destination. direct=True bypasses cooperative partitioning and issues one exact 4-, 8-, or 16-byte copy; that mode requires thread scope and a constant-size region. Only the exact string "zero" changes the operand form; other fill_mode values follow the ordinary predicated or unpredicated branch. The implementation converts prefetch_size to an integer and constructs the PTX spelling, whose table lookup accepts the registered qualifiers.

Demonstration program#

A CTA (128 threads) asynchronously loads a 128×32 float16 tile global → shared, then commits and waits before reading it back (from test_ldgsts.py):

shape, dtype = (128, 32), "float16"
s_layout = TileLayout(S[shape]); full = (slice(0, 128), slice(0, 32))

@Tx.prim_func
def copy_async(A_ptr: Tx.handle, B_ptr: Tx.handle):
    A = Tx.match_buffer(A_ptr, shape, dtype)
    B = Tx.match_buffer(B_ptr, shape, dtype)
    Tx.device_entry(); Tx.cta_id([1]); Tx.warp_id([4]); Tx.lane_id([32]); tid = Tx.thread_id([128])
    A_smem = Tx.alloc_buffer(shape, dtype, scope="shared", layout=s_layout)
    Tx.tile.cta.copy_async(A_smem[full], A[full], dispatch="ldgsts")   # async global -> shared
    Tx.ptx.cp.async_.commit_group()                               # caller commits ...
    Tx.ptx.cp.async_.wait_group(0)                                # ... and waits
    Tx.cuda.cta_sync()
    Tx.tile.cta.copy(B[full], A_smem[full])

Algorithm#

1. Same partition as copy → vec_auto global/shared path. align_layouts_gs builds the [outer, threads, vec] split with the global side driving the canonical order — but the vector candidates are clamped to {128, 64, 32} bits so the byte size is a legal cp.async cp_size. For 128×32 = 4096 float16 over 128 threads the widest legal width is vec = 8 (8 × 2 B = 16 B), giving outer = 4.

2. Emit cp.async instead of a synchronous copy, and do not sync:

for f in Tx.unroll(total_outer):
    s_lin = s_p.apply(f, tid, v0, shape=apply_shape)["m"]
    g_lin = g_p.apply(f, tid, v0, shape=apply_shape)["m"]
    s_ptr = _ptr_off(s_buf.ptr_to(s_zero), _s_off(f, s_lin))
    g_ptr = _ptr_off(g_buf.ptr_to(g_zero), g_lin)
    Tx.ptx[chain](s_ptr, g_ptr, cp_size)   # async; chain includes cache/prefetch modifiers
# NO cta_sync — commit_group / wait_group / cta_sync are the caller's job

Completion is the caller’s responsibility (Tx.ptx.cp.async_.commit_group() then Tx.ptx.cp.async_.wait_group(0)); the dispatch only issues the in-flight loads.

Generated TIRx IR#

for f in Tx.unroll(4):                                   # outer = 4
    s_ptr = pointer_offset(A_smem, ...)
    g_ptr = pointer_offset(A_1, ...)
    Tx.ptx["cp.async.cg.shared.global"](s_ptr, g_ptr, 16)  # cp_size = 16 B

Generated CUDA#

// Core instructions (the generated source wraps these in inline helpers):
asm volatile("cp.async.cg.shared.global [%0], [%1], 16;" ...);  // x4 per thread
asm volatile("cp.async.commit_group;");
asm volatile("cp.async.wait_group 0;");

Each thread issues four asynchronous 16-byte copies; nothing blocks until the caller’s wait_group.

How inputs change the algorithm#

The dtype/alignment set vec (hence cp_size and outer), but unlike the synchronous variant the width is capped at 16 B (cp.async maximum):

case	`vec`	`cp_size`	`outer = 4096 / (128 · vec)`
`float16`, aligned	8	16 B	4
`float32`, aligned	4	16 B	8
8-B-aligned only	(clamped)	8 B	(doubles)
4-B-aligned only	(clamped)	4 B	(×4)

If the region can’t satisfy even a 4-byte (32-bit) cp_size, align_layouts_gs finds no candidate and the variant declines. The direction is fixed: a shared → global copy_async is never ldgsts (hardware has no store form).

copy_async → ldgsts

Contents