Tile Operands and Execution#

A CUDA tile operation repeats a chosen computation over a tensor region. T.cuda.tile.apply uses the PTX instruction you specify, while T.cuda.tile.ewise applies a scalar TIRx function to each element. The operand views describe which values each thread handles and how many times it repeats that computation.

Reading an Operand’s Shape#

Suppose each of 64 threads processes twelve float32 values. A ld.global.v4.b32 instruction loads four 32-bit words at a time, so each thread needs three loads. We can describe this work with shape (2, 32, 3, 4):

(2, 32, 3, 4)
 └───┘  │  └─ atom: four words used by one instruction
  group └──── repeat: three instructions per thread

Part	Meaning in this example
Group	`(2, 32)` selects one of two warps and a lane within that warp.
Repeat	`(3,)` repeats the chosen instruction three times for that thread.
Atom	`(4,)` selects the four words transferred by `.v4.b32`.

The instruction fixes the atom shape. The leading axes describe its repetitions and the threads that execute them. The register layout below maps the first two coordinates to warpid and laneid, leaving 3 * 4 = 12 local positions per thread.

A Load, Compute, and Store Example#

The kernel uses this layout to load the values, add one to each, and store the results. scope="cta" tells the tile calls to use the leading (2, 32) axes for the threads in the block:

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

register_layout = TileLayout(S[(2, 32, 3, 4) : (1 @ warpid, 1 @ laneid, 4, 1)])

@T.function(private=True)
def plus_one(value: T.float32) -> T.float32:
    return value + T.float32(1)

@T.function
def add_one_tile(
    A: T.Tensor((2, 32, 3, 4), "float32", align=16),
    B: T.Tensor((2, 32, 3, 4), "float32", align=16),
):
    T.device_entry(launch=T.cuda.LaunchConfig(grid=1, block=64))
    registers = T.alloc_local((2, 32, 3, 4), "float32", layout=register_layout)
    T.cuda.tile.apply(T.ptx["ld.global.v4.b32"], registers, A, scope="cta")
    T.cuda.tile.ewise(plus_one, registers, registers, scope="cta")
    T.cuda.tile.apply(T.ptx["st.global.v4.b32"], B, registers, scope="cta")

The register layout assigns each thread twelve storage positions: three groups of four. The load and store each transfer four 32-bit words per instruction, so each thread issues three loads and three stores. Between them, ewise applies plus_one to all twelve scalar values.

We can compile and check the result:

import torch
import tvm_tirx

compiled = tvm_tirx.compile(add_one_tile, target="cuda")
a = torch.arange(768, dtype=torch.float32, device="cuda").reshape(2, 32, 3, 4)
b = torch.empty_like(a)
compiled(a, b)
torch.testing.assert_close(b, a + 1)

Execution Scope and Scalar Functions#

The kernel’s LaunchConfig launches 64 threads; scope="cta" describes how the tile operands map to those threads. It does not launch additional threads or insert a barrier. At scope="thread", an operand has no group prefix. Other scopes and warp-collective instructions use the group shapes listed in the coordinate reference.

ewise uses one scalar value per invocation, so it has no trailing instruction atom. For the same register view, (2, 32) still selects a thread, while (3, 4) now specifies twelve scalar invocations.

Preparing Operand Views#

The load’s source and destination expose the same group and repeat coordinates. Within each atom, the four global words must be consecutive and 16-byte aligned; the destination words must belong to the selected thread. These requirements come from the chosen instruction, and lowering checks that the input and output views meet them.

An existing tensor may need a view or rearrange to expose these axes, or a sub to select a region. These operations reuse its storage. They can change the coordinates used to access values, but a transfer between threads or memory spaces still requires an explicit load, store, or other communication instruction. See Types, Tensors, and Memory for tensor views and Data Layout for their physical placement.

The dtype describes storage elements. In the example, one element is one 32-bit word. An instruction that packs two float16 values into one 32-bit register instead needs an appropriate packed container; apply does not pack values automatically. The instruction reference lists the required operand order, representation, alignment, and optional parameters for each family.

Synchronization Belongs to the Caller#

The example uses thread-private registers, so no thread waits for another. If a thread instead reads shared memory written by other threads, the kernel must synchronize those accesses. apply and ewise do not add that synchronization.

Asynchronous instructions also require their own completion protocol. cp.async uses copy groups and waits; TMA loads use transaction barriers; tcgen05 operations have their own commit and wait instructions. The corresponding instruction guides show these calls alongside the data movement. Selecting an execution scope does not perform thread election or complete an asynchronous transfer.

Tile Operands and Execution

Contents