ld / st#
Use T.cuda.tile.apply to repeat a chosen memory instruction. The operand
order is always destination first: a load takes registers, memory and a
store takes memory, registers. For example:
T.cuda.tile.apply(T.ptx["ld.global.v4.b32"], R, A)
T.cuda.tile.apply(T.ptx["st.shared.v4.b32"], S, R)
Both calls transfer four 32-bit words per repeat. Global-to-shared copying through registers consists of these two explicit calls and caller-allocated register storage. To issue a direct asynchronous transfer, use cp.async.
Contract and instruction selection#
Operands have shape (*group,*repeat,word). The final extent is the vector
length in the opcode, or one for a scalar instruction. Both operands use
scalar elements with the opcode’s storage bit width. Each memory atom must
be consecutive and aligned to the whole vector’s byte size. Its start can
vary across repeats. Local register slots can be separated or permuted, but
each must belong to the issuing thread; loaded destinations cannot overlap.
The supported address-space spellings are global, shared,
shared::cta, and local. The pointer side of ld.local/st.local
is addressable local memory, distinct from the register operands. Legal
types, vector lengths, and qualifiers come from the selected PTX form.
Cache operators such as .nc and .cg and prefetch/eviction qualifiers
are part of the opcode. A form with a cache-policy operand takes
cache_policy=.... pred gates execution. A predicated load can use
preserve_dst=True to keep an already initialized destination when the
predicate is false. Lowering does not reduce the width when alignment fails.
Example: a strided tensor and gapped registers#
The following kernel reads a float32 tensor whose logical axes are not in physical order. It keeps that layout and prepares a view with four physically adjacent words at the end. The destination deliberately has twice the storage stride, so each vector writes separated local slots.
from tvm.script import tirx as T
from tvm.tirx.layout import S, TileLayout
global_layout = TileLayout(S[(15, 2, 7, 2) : (4, 1, 60, 2)])
register_layout = TileLayout(S[(15, 2, 7, 2) : (8, 2, 120, 4)])
@T.function
def strided_load(
A: T.Tensor((15, 2, 7, 2), "float32", layout=global_layout, align=16),
Out: T.Tensor((420,), "float32"),
):
T.device_entry(launch=T.cuda.LaunchConfig(grid=1, block=1))
R = T.alloc_tensor(
ty_args=[T.Tensor((15, 2, 7, 2), "float32", scope="local", layout=register_layout)]
)
av = A.rearrange("i b j d -> j i (d b)")
rv = R.rearrange("i b j d -> j i (d b)")
T.cuda.tile.apply(T.ptx["ld.global.v4.b32"], rv, av)
for j, i, d, b in T.grid(7, 15, 2, 2):
Out[60 * j + 4 * i + 2 * d + b] = R[i, b, j, d]
Both views have shape (7,15,4). For one repeat (j,i), source word
w lives at 60*j+4*i+w and destination word w lives at local
offset 120*j+8*i+2*w. One instruction binds its four outputs directly to
these four slots. This emits 105 ld.global.v4.b32 instructions, with no
intermediate tile or register-remapping copy. The final scalar stores are
ordinary caller code used to inspect the result.
Compile and check the complete kernel:
import numpy as np
import tvm
compiled = tvm.tirx.compile(strided_load, target="cuda", pipeline="tirx")
x = np.arange(420, dtype="float32")
a = tvm.runtime.tensor(x.reshape(15, 2, 7, 2), tvm.cuda())
out = tvm.runtime.tensor(np.zeros_like(x), tvm.cuda())
compiled(a, out)
np.testing.assert_array_equal(out.numpy(), x)
The raw load in the generated PTX has four output operands, regardless of the local offsets selected by the view:
ld.global.v4.b32 {r0, r1, r2, r3}, [address];