TMA Transfers

This guide shows how to use the Tensor Memory Accelerator (TMA) to move multidimensional tiles between global and shared memory. One thread issues the transfer using a TensorMap descriptor, which records the global shape, strides, element format, and tile configuration. TIRx prepares this descriptor in the host launcher from the tensor operand and explicit box parameters.

TMA requires a Hopper-or-later GPU. The Async Data Movement: TMA chapter explains its hardware role and layout support in detail.

A Tile Load, Computation, and Store

The following kernel loads an (8, 32) float32 tile, adds one to each element, and stores it back to global memory. Lane 0 issues the transfers; all 32 lanes update one column each:

from tvm_tirx.script import tirx as T

@T.function
def tma_add_one(
    A: T.Tensor((8, 32), "float32", align=16),
    B: T.Tensor((8, 32), "float32", align=16),
):
    T.device_entry(launch=T.cuda.LaunchConfig(grid=1, block=32))
    lane = T.cuda.lane_id()
    shared = T.alloc_shared((8, 32), "float32", layout="default", align=128)
    mbar = T.alloc_shared((1,), "uint64", align=8)
    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.ptx.fence.proxy.async_.shared__cta()
    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, 32), coords=(0, 0), mbar=mbar.ptr_to([0]),
        )
    T.cuda.mbarrier_wait(mbar.ptr_to([0]), 0)
    for row in range(8):
        shared[row, lane] = shared[row, lane] + 1
    T.ptx.fence.proxy.async_.shared__cta()
    T.cuda.cta_sync()
    if lane == 0:
        T.cuda.tile.apply(
            T.ptx["cp.async.bulk.tensor.2d.global.shared::cta.tile.bulk_group"],
            B, shared, box=(8, 32), coords=(0, 0),
        )
        T.ptx.cp.async_.bulk.commit_group()
        T.ptx.cp.async_.bulk.wait_group(0)
    T.cuda.cta_sync()

box=(8, 32) selects the tile dimensions and coords=(0, 0) selects its starting position in the global tensor, both in outer-to-inner order. The shared tensor has the same tile shape. layout="default" gives it contiguous row-major storage, and align=128 aligns its base. The load takes shared, A; the store reverses the direction with B, shared.

Load Completion

The load reports completion through mbar, an eight-byte barrier object in shared memory. Its initialization expects one arrival per phase. arrive.expect_tx performs that arrival and registers the tile’s 8 * 32 * 4 = 1024 transaction bytes. The initial fence and CTA barrier publish the initialized barrier before lane 0 issues the transfer.

The TMA engine completes those transaction bytes as it fills shared memory. mbarrier_wait(..., 0) waits for the initial phase to finish, so each lane can then read and update its column. A phase completes when both its expected arrivals and transaction bytes have been satisfied. The mbarrier chapter explains phases and reuse across pipeline stages.

Store Completion

After the updates, each thread executes the proxy fence to make its shared-memory writes visible to the asynchronous transfer. The CTA barrier then ensures all writers have finished their writes and fences before lane 0 issues the store.

A TMA store uses a bulk async group. Lane 0 commits its store with bulk.commit_group() and waits for completion with bulk.wait_group(0). The final CTA barrier holds the other threads until that wait finishes, allowing the shared buffer to be reused. These bulk groups are distinct from the ordinary copy groups used by Asynchronous Copies.

Compile and Run

The host launcher needs tensor-map encoding support, available with an LLVM-enabled build or the CUDA C++ host backend. With LLVM enabled, the kernel can be compiled and checked as follows:

import torch
import tvm_tirx

a = torch.arange(256, dtype=torch.float32, device="cuda").reshape(8, 32)
b = torch.empty_like(a)
tvm_tirx.compile(tma_add_one, target="cuda")(a, b)
torch.testing.assert_close(b, a + 1)

The TMA instruction reference covers repeated boxes, coordinate steps, swizzles, multicast, reductions, and packed formats. For a complete multistage schedule, see Pipelining GEMM with TMA.