Data Movement

This guide shows how to move tensor data between global memory, shared memory, and thread-local registers. The instruction determines the source and destination spaces, while tensor views select the values to transfer. All three instruction families take the destination first:

Instruction	Operands	Data movement
ld	registers, memory	Memory to the executing thread’s registers
st	memory, registers	Registers to memory
mov	destination, source	Register values within the same thread

A Shared-Memory Exchange

A block can use shared memory to exchange data between threads. In this example, each of 64 threads writes its four-element row, waits for the other threads, then reads the next thread’s row:

from tvm_tirx.script import tirx as T

@T.function
def 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")
    registers = T.alloc_local((4,), "float32", layout="default")
    shared = T.alloc_shared((64, 4), "float32", align=16)
    T.cuda.tile.apply(T.ptx["ld.global.v4.b32"], registers, A.sub[tid, :])
    T.cuda.tile.apply(T.ptx["st.shared.v4.b32"], shared.sub[tid, :], registers)
    T.cuda.cta_sync()
    T.cuda.tile.apply(
        T.ptx["ld.shared.v4.b32"], registers, shared.sub[(tid + 1) % 64, :]
    )
    T.cuda.tile.apply(T.ptx["st.global.v4.b32"], B.sub[tid, :], registers)

A.sub[tid, :] selects a four-element row without reading it. The first apply performs the read; the next writes those values to shared memory. The barrier makes every thread’s write visible before the shared load. Thread 63 reads row 0, so the output shifts the rows cyclically.

Each .v4.b32 instruction transfers four 32-bit words. The row views and register tensor have shape (4,): one four-word atom with no repeat axes. The calls use the default thread scope, so each thread transfers its own row. The 16-byte alignment and four-element row stride keep every memory atom aligned.

We can compile the kernel and compare its result with a row shift:

import torch
import tvm_tirx

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

The two global/shared transfers pass through the caller’s register tensor. For a direct asynchronous global-to-shared transfer, see Asynchronous Copies, which includes the required completion calls.

Register Copies and Initialization

Inside this kernel, we can save loaded values before reusing registers for another read. Insert the following statements after its global load:

saved = T.alloc_local((4, 1), "float32", layout="default")
T.cuda.tile.apply(T.ptx["mov.b32"], saved, registers.view(4, 1))

The view exposes four repeats, each ending in a single-register atom (1,). The call copies the four words into saved. Later operations can use that copy after registers has been overwritten.

A scalar source broadcasts to every repeat. For example, the following call initializes all four saved values to zero:

T.cuda.tile.apply(T.ptx["mov.b32"], saved, T.float32(0))

mov preserves the bits within the same thread. A numeric type change uses cvt, covered in Arithmetic and Type Conversion. Transfers between threads use shared memory or another communication instruction.

Cache qualifiers, predication, and destination preservation are described in the memory instruction reference and register move reference.