cp.async.bulk.tensor / cp.reduce.async.bulk.tensor#

TMA transfers a descriptor-defined global box to or from a shared-memory image. Select direction, rank, tile/gather/scatter mode, multicast, cache qualifiers, and any reduction operator in the opcode. apply encodes the specified descriptor and issues that instruction at explicit coordinates; it does not plan a decomposition of an arbitrary tensor layout.

Operands and descriptor fields#

For a load the operands are shared, global; for a store/reduction they are global, shared. The global operand is either a rectangular global tensor region or an ordinary pointer to a caller-encoded TensorMap. In the latter case the caller is responsible for matching its fields to the tile contract; lowering does not inspect or select opaque descriptors.

Keyword	Contract
`box`	Required extents in outer-to-inner order; one per instruction dimension, each 1..256.
`coords`	Required runtime starting coordinates in the same order, relative to the descriptor view.
`coord_steps`	One rank-length increment vector per shared group/repeat axis; required when such axes exist.
`format`	Storage dtype by default, or an explicit matching format such as `tf32` over float32.
`swizzle`	0/1/2/3: none/32B/64B/128B shared swizzle.
`element_strides`	Traversal steps in 1..8; default all ones. Without interleave, the innermost step is one.
`interleave`	0/1/2: none/16B/32B. The global and shared images change as described below.
`l2_promotion`, `oob_fill`	Explicit descriptor policies; both default to zero.
`mbar` / `cta_mask` / `cache_policy`	Runtime operands required by load / multicast / cache-hint forms, respectively.

Without interleave, the global region has exactly the opcode rank, unit innermost stride, and suitably aligned outer byte strides. Its dimensions and outer strides can be runtime values and are passed to the host descriptor encoder. Known invalid fields are rejected rather than silently promoted, folded, transposed or reinterpreted as another dtype.

The shared suffix is the box image: ceil(box[i]/element_strides[i]) per dimension, compact before swizzling. Every atom’s unswizzled origin is 128-byte aligned, even for swizzle zero; a swizzled backing base is aligned to the whole swizzle period. Swizzled subregions use their inverse-swizzled origins so the hardware applies the phase once.

Example: two explicit 8x16 boxes#

from tvm.script import tirx as T


@T.function
def tma_load(A: T.Tensor((16, 16), "uint32"), Out: T.Tensor((16, 16), "uint32")):
    T.device_entry(launch=T.cuda.LaunchConfig(grid=1, block=32))
    lane = T.cuda.lane_id()
    shared = T.alloc_shared((2, 8, 16), "uint32", layout="default", align=128)
    mbar = T.alloc_shared((1,), "uint64", layout="default")
    if lane == 0:
        T.ptx.mbarrier.init.shared.b64(mbar.ptr_to([0]), T.uint32(1))
        T.ptx.mbarrier.arrive.expect_tx.shared.b64(mbar.ptr_to([0]), T.uint32(1024))
    T.cuda.cta_sync()
    if lane == 0:
        T.cuda.tile.apply(
            T.ptx["cp.async.bulk.tensor.2d.shared::cta.global.tile.mbarrier::complete_tx::bytes"],
            shared,
            A,
            box=(8, 16),
            coords=(0, 0),
            coord_steps=((8, 0),),
            mbar=mbar.ptr_to([0]),
        )
    T.cuda.mbarrier_wait(mbar.ptr_to([0]), 0)
    T.cuda.cta_sync()
    for g, r in T.grid(2, 8):
        if lane < 16:
            Out[g * 8 + r, lane] = shared[g, r, lane]

The shared shape is (repeat=2,box_row=8,box_col=16). One elected lane issues two instructions at global coordinates (0,0) and (8,0). Each copies 512 bytes; one phase of the barrier accounts for 1024 bytes. Descriptor construction is host initialization associated with this call. It is not repeated on the device for each box.

import numpy as np
import tvm

x = np.arange(256, dtype="uint32").reshape(16, 16)
compiled = tvm.tirx.compile(tma_load, target="cuda", pipeline="tirx")
out = tvm.runtime.tensor(np.zeros_like(x), tvm.cuda())
compiled(tvm.runtime.tensor(x, tvm.cuda()), out)
np.testing.assert_array_equal(out.numpy(), x)

Stores, reductions, and completion#

The corresponding store spelling is cp.async.bulk.tensor.2d.global.shared::cta.tile.bulk_group. Reduction uses, for example, cp.reduce.async.bulk.tensor.2d.global.shared::cta.add.bulk_group. The coordinate/box contract is the same, with global destination first. Complete shared initialization and issue T.ptx.fence.proxy.async_.shared__cta() before the transfer; the issuing thread then calls T.ptx.cp.async_.bulk.commit_group() and T.ptx.cp.async_.bulk.wait_group(0). Initialize a reduction destination to the intended starting value. No mbarrier operand is used for these stores.

Loads use their supplied mbarrier, with explicit initialization, expected transaction bytes, and phase wait. Multicast adds the explicit CTA mask and the corresponding cluster/barrier setup. The tile call elects no thread and inserts neither waits nor descriptor prefetch. A desired prefetch is an ordinary T.ptx.prefetch.tensormap call.

Gather4 and scatter4#

These forms use a 2D descriptor with box=(1,width) and a shared image (4,width). The four row coordinates are ordinary runtime operands in gather4; they need not be consecutive. For a prepared shared of shape (4,32) and a matching barrier, an issue can be written as:

T.cuda.tile.apply(
    T.ptx["cp.async.bulk.tensor.2d.shared::cta.global.tile::gather4.mbarrier::complete_tx::bytes"],
    shared, A, box=(1,32), coords=(0,0),
    gather4=(row,row+3,row+7,row+11), mbar=mbar.ptr_to([0]),
)

For uint32 this is one 512-byte gather4 issue. Scatter4 uses global destination first and the matching store opcode/completion protocol. Gather/scatter coordinate selection is caller logic.

Interleave and packed FP4#

With interleave=1 or 2, a rank-3..5 opcode takes global axes (N,C_slices,*spatial,Ci). Ci holds one 16B/32B slice and is implicit in the descriptor rank. Its elements are contiguous and the innermost spatial stride is one slice. Interleave 2 requires swizzle 1 and 32-byte global alignment. One issue selects one C slice, even if a larger C box field is supplied; channel iteration requires an explicit repeat and C step.

For example, uint16/interleave-2 with box=(1,1,8) has shared atom (1,1,8,16). Add a leading repeat of three and coord_steps=((0,1,0),) to issue three successive channel slices. Gather4/scatter4 do not use interleave. On B200, 16B interleave with swizzle 1/2/3 can trap, including when issued through raw PTX.

format="float4_e2m1fn" selects the load-only 16U4_ALIGN16B format. The inner global dimension is a multiple of 128, the inner box is 128, and global origins/outer strides are 32-byte aligned. Each sixteen nibbles occupy the low eight bytes of a sixteen-byte shared block. The remaining eight bytes are undefined padding; barrier accounting counts payload bytes, not this padding. Swizzle 0/3 and gather4 use this explicit packed image.

Implementation: python/tvm/backend/cuda/tile/patterns/tma.py.

cp.async.bulk.tensor / cp.reduce.async.bulk.tensor

Contents