ldmatrix / stmatrix

ldmatrix / stmatrix#

The tile patterns support m8n8.b16 with .x1, .x2, or .x4 and optional .trans. Select the exact form in the opcode:

T.cuda.tile.apply(
    T.ptx["ldmatrix.sync.aligned.m8n8.x2.shared.b16"],
    registers, shared_words, scope="warp",
)

stmatrix reverses the operand order: shared_words, registers. The implementation repeats the specified warp instruction, without picking a matrix count or transpose mode from the layout.

Address providers and result owners#

Both regions end in (matrix,row,word)=(count,8,4) and use scalar 32-bit containers. A container holds two adjacent 16-bit values. Use a dtype view to expose packed half/bfloat16 storage as b32 words; this is not a conversion.

  • Each shared row has four consecutive words and a 16-byte-aligned start. Different rows may have gaps and may be interleaved with other matrices.

  • Lane 8*matrix+row provides that row’s shared address. For x1/x2, the remaining address providers repeat valid row addresses.

  • Register word (matrix,row,word) belongs to lane 4*row+word. Each lane receives count words. Their register offsets need not be adjacent, and need not have the same order as another tensor declaration.

Thus, the lane supplying a shared address generally differs from the lanes receiving the values from that row. The contract checks both maps. All 32 lanes must participate, including those whose addresses are duplicates.

Example: interleaved shared rows and register slots 0, 2, 8, 10#

Here the original shared tensor is (16,4,2):(8,1,128). The first matrix uses even original rows, the second uses odd rows, and the final original axis selects the second instruction. The factory below includes shared initialization, synchronization, and readback for one warp.

from tvm.script import tirx as T
from tvm.backend.cuda.axis import laneid
from tvm.tirx.layout import S, TileLayout

LDMATRIX = "ldmatrix.sync.aligned.m8n8.x2.shared.b16"

def matrix_kernel():
    shared = TileLayout(S[(16, 4, 2) : (8, 1, 128)])
    registers = TileLayout(S[(2, 2, 8, 4) : (8, 2, 4 @ laneid, 1 @ laneid)])

    @T.function
    def kernel(A: T.Tensor((16, 4, 2), "uint32"), Out: T.Tensor((32, 12), "uint32")):
        T.device_entry(launch=T.cuda.LaunchConfig(grid=1, block=32))
        lane = T.cuda.lane_id()
        SM = T.alloc_tensor(
            ty_args=[T.Tensor((16, 4, 2), "uint32", scope="shared", layout=shared, align=16)]
        )
        for q in range(4):
            n = lane + 32 * q
            SM[n // 8, (n // 2) % 4, n % 2] = A[n // 8, (n // 2) % 4, n % 2]
        T.cuda.cta_sync()
        raw = T.alloc_tensor(ty_args=[T.Tensor((12,), "uint32", scope="local", layout="default")])
        for i in T.unroll(12):
            raw[i] = T.uint32(4294966519)  # bit pattern of int32(-777)
        R = T.decl_tensor(
            raw.ptr_to([0]),
            ty_args=[T.Tensor((2, 2, 8, 4), "uint32", scope="local", layout=registers)],
        )
        S_std = SM.rearrange("(r m) w k -> k m r w", r=8, m=2)
        T.cuda.tile.apply(T.ptx[LDMATRIX], R, S_std, scope="warp")
        for i in T.unroll(12):
            Out[lane, i] = raw[i]

    return kernel

S_std[k,m,r,w] is SM[2*r+m,w,k] and has strides (128,8,16,1). Both operands have shape (2,2,8,4): one repeat axis of length two, then the x2 matrix atom. For lane l the loaded values are:

raw[0]  = SM[2*(l//4),   l%4, 0]
raw[2]  = SM[2*(l//4)+1, l%4, 0]
raw[8]  = SM[2*(l//4),   l%4, 1]
raw[10] = SM[2*(l//4)+1, l%4, 1]

All other local slots retain their initialized sentinel. The transfer emits two x2 instructions per warp. It does not emit x4 and does not copy from a contiguous temporary fragment into the separated output slots.

Run the complete factory example on an Ampere-or-later device:

import numpy as np
import tvm

compiled = tvm.tirx.compile(matrix_kernel(), target="cuda", pipeline="tirx")
x = np.arange(128, dtype="uint32").reshape(16, 4, 2)
out = tvm.runtime.tensor(np.zeros((32, 12), dtype="uint32"), tvm.cuda())
compiled(tvm.runtime.tensor(x, tvm.cuda()), out)
expected = np.full((32, 12), 4294966519, dtype="uint32")
for lane in range(32):
    for k in range(2):
        for matrix in range(2):
            expected[lane, 8 * k + 2 * matrix] = x[2 * (lane // 4) + matrix, lane % 4, k]
np.testing.assert_array_equal(out.numpy(), expected)

Transpose and stores#

.trans transposes each 8x8 matrix of 16-bit elements, not a matrix of b32 words. Register image element (row,2*word+half) comes from shared element (2*word+half,row). The register word’s lane owner remains 4*row+word. Reorganize operands with this image in mind; toggling .trans is a semantic change, not a layout optimization hint.

stmatrix uses the corresponding register image as input and each provider’s shared row address as output. Its supported architecture is checked by the PTX compiler. Synchronize before a load consumes another thread’s shared stores, and after a store before other threads consume or reuse that shared storage. No fence or barrier is inserted by the pattern.

An insufficiently aligned row, wrong register owner, overlapping loaded registers, or a partially active warp is rejected. Gaps between rows or register slots are legal; gaps within a shared 16-byte row are not.

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

Contents