Layout IR#

Layouts describe how logical tensor coordinates map to storage and execution axes. See Tensor Layout for the programming model and worked examples.

The primary classes are re-exported from tvm.tirx, but are defined in this module:

class tvm.tirx.layout.Layout

verify_well_formed() → bool

Verify if the layout is well-formed.

Returns:: True if the layout is well-formed, False otherwise
Return type:: bool

size(axis_name: str | None = None)

Get the size of the layout.

Parameters:: axis_name (Optional[str]) – The name of the axis to get the size of. If not provided, the default input size will be returned.

span(axis_name: str | None = None)

Get the span of the layout.

Parameters:: axis_name (Optional[str]) – The name of the axis to get the span of. If not provided, the default span will be returned.

apply(*coord: list[Expr], shape: list[Expr] | None = None) → dict[str, Expr]

Apply the layout on the input coordinate and get the mapped output.

Input cases: - coord is a single element -> will be treated as a 1D coordinate - coord is a list of elements -> will be treated as a multi-dimensional coordinate - shape is provided -> turn the coord with shape into a 1D coordinate - shape is not provided -> use the default shape

Returns:: The mapped output (axis name -> value on the axis)
Return type:: Dict[str, Expr]

apply_to_shape(coord: list[Expr], input_shape: list[Expr]) → list[Expr]

Compute the per-shard value that each shard would take if coord were interpreted against input_shape.

Tries self.group(input_shape) first. On success, each group owns exactly one input_shape entry, so coord[d] can be split within that group’s shard extents (bounds stay local to one input dim — simpler analyzer simplification, no cross-dim complications).

Falls back to FlattenCoord(coord, input_shape) + SplitCoord on self’s raw shard shape when the group call fails (e.g. when input_shape does not align with the layout’s factor boundaries).

Returns a list of length len(self.shard); each entry is the value that shard would iterate.

canonicalize() → Layout

Canonicalize the layout by simplifying and fusing iterators where possible.

Returns:: The canonicalized layout
Return type:: Layout

tile(outer: TileLayout, outer_shape: list[Expr], inner_shape: list[Expr]) → TileLayout | ComposeLayout

Tile the current layout with an outer layout.

Parameters:

outer (TileLayout) – The outer layout to tile with
outer_shape (List[Expr]) – The shape of the outer layout
inner_shape (List[Expr]) – The shape of the inner layout

Returns:

The resulting tiled layout

Return type:

Union[TileLayout, ComposeLayout]

direct_sum(left: TileLayout, left_shape: list[Expr], right_shape: list[Expr]) → TileLayout | ComposeLayout

Direct-sum on the tiling domain (unscaled composition): A + B.

This layout is treated as the right addend B grouped by right_shape. The left layout is treated as A grouped by left_shape. The resulting layout is evaluated over the interleaved domain S_A ⊗ S_B, without span scaling (unlike tiling).

is_tile_inner(tile_layout: TileLayout | ComposeLayout, tiled_shape: list[Expr], inner_shape: list[Expr]) → TileLayout | None

Check if a layout is the inner layout of a tiled layout.

Parameters:

tile_layout (Union[TileLayout, ComposeLayout]) – The tiled layout to check
tiled_shape (List[Expr]) – The shape of the tiled layout
inner_shape (List[Expr]) – The shape of the inner layout

Returns:

The outer layout if it is the inner layout of the tiled layout, None otherwise

Return type:

Optional[TileLayout]

is_tile_outer(tile_layout: TileLayout | ComposeLayout, tiled_shape: list[Expr], outer_shape: list[Expr]) → Layout | None

Check if a layout is the outer layout of a tiled layout.

Parameters:

tile_layout (Union[TileLayout, ComposeLayout]) – The tiled layout to check
tiled_shape (List[Expr]) – The shape of the tiled layout
outer_shape (List[Expr]) – The shape of the outer layout

Returns:

The inner layout if it is the outer layout of the tiled layout, None otherwise

Return type:

Optional[Layout]

is_direct_sum_right(sum_layout: TileLayout | ComposeLayout, interleaved_shape: list[Expr], right_shape: list[Expr]) → TileLayout | None

Check if this layout is the right addend B in a direct-sum A + B.

Returns the left addend A if recognized, otherwise None.

is_direct_sum_left(sum_layout: TileLayout | ComposeLayout, interleaved_shape: list[Expr], left_shape: list[Expr]) → Layout | None

Check if this layout is the left addend A in a direct-sum A + B.

Returns the right addend B if recognized, otherwise None.

slice(shape: list[Expr], region: list[tuple[Expr, Expr]]) → Layout | None

Slice the layout with a given shape and region.

Parameters:

shape (List[Expr]) – The shape of the layout
region (List[Tuple[Expr, Expr], tvm.ir.Range]) – The region to slice, each element is (begin, end)

Returns:

The sliced layout, or None if slicing is not possible

Return type:

Optional[Layout]

tile_to(to_shape: list[Expr], current_shape: list[Expr]) → Layout

Tile the current layout to the given shape.

Parameters:

to_shape (List[Expr]) – The shape to tile to
current_shape (List[Expr]) – The current shape of the layout

is_swizzle() → bool: Check if the layout is a bare swizzle (ComposeLayout over a trivial tile).

is_trivial() → bool: Check if the layout is trivial.

is_trainium() → bool: Check if the layout is trainium layout.

unpack(num: int) → Layout

Unpack the layout, where a single element in the layout is unpacked into num contiguous elements.

Parameters:: num (int) – The number of elements to unpack into
Returns:: The unpacked layout
Return type:: Layout

broadcast(num: int, position: int = -1, axis: 'Axis' | str = 'm') → Layout

Insert a stride-0 broadcast dim of extent num at position.

position follows Python list-insert semantics (negative indices count from the end; -1 appends after the last shard dim). The new dim has stride 0 — accessing along it doesn’t move the byte offset, so the same physical element is “seen” num times.

Useful for layouts where a consumer reads the same SMEM datum multiple times (e.g. sf_reuse over MMA-K steps).

pack(num: int) → Layout

Pack the layout, where num contiguous elements in the layout are packed into a single element.

Parameters:: num (int) – The number of elements to pack into
Returns:: The packed layout
Return type:: Layout

class tvm.tirx.layout.Axis(*args, **kwargs)

Layout axis wrapper.

static get(name: str) → Axis: Get or create an axis by name. Unknown names are auto-registered.

is_thread() → bool: Check if the axis is a thread axis.

is_memory() → bool: Check if the axis is a memory axis.

get_scope() → ExecScope | None: Get the scope of the axis.

get_subscope() → ExecScope | None: Get the subscope of the axis.

class tvm.tirx.layout.Iter(extent: Expr, stride: Expr, axis: Axis | str): A memory layout that tiles data across devices.

class tvm.tirx.layout.TileLayout(spec: _LayoutSpec)

A memory layout that tiles data across devices.

static from_iters(shard: Sequence[Iter] = (), replica: Sequence[Iter] = (), offset: dict[Axis | str, Expr] | None = None) → TileLayout: Construct a TileLayout from pre-built Iter objects.

is_trivial() → bool: Check if the layout is trivial.

group(shape: list[Expr]) → tuple[Layout, list[int]]

Group the current layout by the given shape.

Parameters:: shape (List[Expr]) – The shape to group by
Returns:: The grouped layout and the separators
Return type:: Tuple[Layout, List[int]]

group_many(shapes: Sequence[Sequence[Expr]]) → tuple[TileLayout, list[list[int]]]

Group the layout by the minimal common refinement of several shapes.

Repeated cumulative product boundaries are retained, so an extent-one dimension in any input shape becomes a real unit iterator in the refined layout. This operation only splits existing shard iterators; it does not canonicalize or reorder the layout.

Parameters:: shapes (Sequence[Sequence[Expr]]) – Logical shapes with provably equal total products.
Returns:: The commonly refined layout and one separator list per input shape.
Return type:: Tuple[TileLayout, List[List[int]]]

get_scope() → tuple[ExecScope, ExecScope] | None: Get the scope pair of the layout.

classmethod trainium(annotation: str, shape: tuple[Expr], is_psum: bool = False) → TileLayout: Create a TileLayout from an annotation string and a shape.

to_psum() → TileLayout: Convert the layout to a psum layout.

permute_dims(perm: list[int]) → TileLayout: Permute the dimensions of the layout.

permute_by_groups(seps: list[int], perm: list[int]) → TileLayout

Permute groups of shard iters defined by seps.

seps follows the convention of group()’s second return value: seps[0] == 0 and group i covers shard indices [seps[i], seps[i + 1]). The number of groups is len(seps) - 1.

Parameters:

seps (list[int]) – Group boundary positions in the shard list.
perm (list[int]) – Permutation of range(len(seps) - 1) selecting the new group order.

class tvm.tirx.layout.ComposeLayout(per_element: int, swizzle_len: int, atom_len: int, tile_layout: TileLayout, swizzle_inner: bool = True)

A memory layout that swizzles a tile layout.

per_element / swizzle_len / atom_len / swizzle_inner carry the swizzle (formerly the standalone SwizzleLayout); tile_layout is the tiled memory map the swizzle is applied to. A bare swizzle is a ComposeLayout over a trivial identity tile.

S[...] and R[...] build shard and replica layout specifications, respectively. Named axes such as laneid, warpid, tid_in_wg, TLane, and TCol are resolved lazily by this module.

Definition of layout.

tvm.tirx.layout.tcgen05_atom_layout(instr_shape: str, tensor_shape: tuple[int, int], dtype) → TileLayout

Describes the per-warpgroup register tile that Tx.copy_async produces when reading a TMEM fragment via tcgen05.{ld,st}.<instr_shape>.xN. rep (the .xN qualifier) is inferred from tensor_shape.

Fragment row count is determined by instr_shape: .32x32b normally covers an M=128 fragment (128 rows per warpgroup), and .16x{64,128,256}b covers an M=64 fragment (64 rows per warpgroup). ("32x32b", (64, N)) is the fp32 Layout B readback image for a .cta_group::2 M=64 accumulator.

TMEM is kept dense for 16-bit dtypes: two 16-bit elements per 32-bit TMEM cell (matching the existing .32x32b convention). The PTX op is issued with the plain .b32 form (no .pack::16b qualifier), and the returned layout describes the per-thread register file with two packed 16-bit elements per 32-bit register.

Parameters:

instr_shape (str) – The PTX atom’s .shape qualifier. One of "32x32b", "16x64b", "16x128b", "16x256b".
tensor_shape (tuple[int, int]) –
The logical fragment shape in element units. Must be (frag_rows, K) where frag_rows is 128 for .32x32b and 64 for the other shapes. The fp32 Layout B image is the exception: it uses ("32x32b", (64, N)). The column extent is divisible by the per-warp column factor for the chosen instr_shape and dtype:
```
K must be a power-of-two multiple of (factor_fp32 * elem_per_32b)
```
where factor_fp32 is 1 / 2 / 4 / 8 for .32x32b / .16x64b / .16x128b / .16x256b, and elem_per_32b is 1 for fp32 and 2 for fp16/bf16. The inferred rep must be in PTX Table 49’s supported set for the chosen instr_shape.
dtype (str | tvm.DataType) – Element dtype. "float32", "float16", or "bfloat16".

Returns:

A tensor_shape-shaped tile layout. The factory builds it as a sequence of fine-grained iters describing the per-(lane, register) destination position.

Return type:

TileLayout

Examples

tcgen05_atom_layout("16x64b", (64, 64), "float32") → .16x64b.x32 (rep=32, fp32).

tcgen05_atom_layout("16x128b", (64, 256), "float16") → .16x128b.x32 (rep=32, fp16; two fp16 elements packed per 32-bit register and per 32-bit TMEM cell).

tvm.tirx.layout.tmem_datapath_layout(datapath: str, rows: int, cols: int, sub_slab: int = 0) → TileLayout

Return the TileLayout for a tcgen05 MMA datapath.

See PTX ISA §9.7.16.10.5 for the datapath enumeration. The returned layout is shape-compatible with a buffer of (rows, cols) and encodes the logical-row → physical-TMEM-lane mapping that the corresponding MMA writes to (and that the matching .16x*b / .32x32b atom expects to read).

Parameters:

datapath (str) –
PTX data path layout letter (all cta_group::2 layouts are the per-CTA view of the M-rows this CTA of the pair owns):
- "A": M=256, .cta_group::2 — identity over 128 rows.
- "B": M=128, .cta_group::2, dense A — 64 rows, column halves folded across the two 64-lane halves.
- "C": M=128, .cta_group::2, sparse A — 64 rows scattered 16-per-warp, half datapath (same organization as F).
- "D": M=128, .cta_group::1 — identity, full datapath.
- "E": M=64, .ws — 64 rows, column halves folded (same organization as B).
- "F": M=64, non-.ws — 64 rows scattered 16-per-warp, half datapath.
- "G": M=32, .ws — 32 rows, column quarters folded across the four 32-lane warp slabs.
rows (int) – Logical row count of the TMEM buffer. Must match the datapath’s M dimension: 128 for A/D, 64 for B/C/E/F, 32 for G.
cols (int) – Logical column count. Datapath B requires an even count because its columns split into two equal lane halves.
sub_slab (int) – For Layout F, select the lower (0) or upper (1) 16-lane half of each warp’s 32-lane TMEM partition. The upper half is useful as a 64-row read/write view of the high half-slab of a Layout D accumulator. Layouts D and B already span both halves and therefore only accept 0.

Returns:

Buffer-shape-compatible layout for (rows, cols).

Return type:

TileLayout

tvm.tirx.layout.tmem_mma_operand_layout(operand, shape, dtype, *, M, cta_group, ws=False, sparse=False, group=None)

Resolve the TMEM TileLayout for a tcgen05 MMA operand from operand role + instruction params.

operand: "A" (A-in-TMEM) or "D" (accumulator / C). M is the PTX tcgen05.mma instruction M (256/128/64), NOT the per-CTA buffer rows; per_cta_rows = M // cta_group. Physical 32-bit column count is left to the pool’s _resolve_cols (layout-aware).

Reproduces the hand-written layouts the three FlashMLA kernels use; the supported/reject domain is spec’d in localdoc/claude_plan.txt S3a.

tvm.tirx.layout.wg_local_layout(cols, rows=128)

Return a warpgroup-local register layout.

The logical (rows, cols) tile is distributed on tid_in_wg along rows, so each thread owns one row and contiguous cols local elements.