Layout IR

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.tirx, but are defined in this module:

class tvm_tirx.tirx.layout.Layout(*args: Any, **kwargs: Any)

Base class for mappings from logical indices to named physical axes.

Concrete layouts describe memory offsets and thread ownership. apply() returns coordinates keyed by axis name; size() describes the input size and span() the covered output span.

check_well_formed() → None

Check whether the layout is well formed.

Raises:

ValueError, RuntimeError – If the layout is not well formed.

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]

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]

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_trivial() → bool

Check if the layout is trivial.

unpack(num: int, axis: str = 'm') → 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

  • axis (str) – Storage coordinate whose element size changes. Ownership coordinates are preserved. Tensor-memory dtype views use TCol.

Returns:

The unpacked layout

Return type:

Layout

pack(num: int, axis: str = 'm') → 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

  • axis (str) – Storage coordinate whose element size changes. Ownership coordinates are preserved. Tensor-memory dtype views use TCol.

Returns:

The packed layout

Return type:

Layout

class tvm_tirx.tirx.layout.Axis(*args: Any, **kwargs: Any)

A named physical coordinate used by a layout.

Use Axis.get(name) or a registered attribute such as Axis.m. Axis metadata identifies memory or thread coordinates and any owning backend; value @ axis attaches a stride or offset to that coordinate.

Examples

from tvm_tirx.script import tirx as T

memory = T.Axis.get("m")  # Also available as T.Axis.m.
layout = T.TileLayout(T.S[8 : 2 @ memory])
assert int(layout.apply(3)["m"]) == 6
property name: str

Return the canonical axis name.

classmethod get(name: str) → Axis

Get or create the axis singleton named name.

Unknown names are registered without thread or memory attributes.

is_thread() → bool

Check if the axis is a thread axis.

is_memory() → bool

Check if the axis is a memory axis.

property backend: str | None

Backend owning this axis, or None for a generic axis.

get_scope() → ThreadScope | None

Get the scope of the axis.

get_subscope() → ThreadScope | None

Get the subscope of the axis.

class tvm_tirx.tirx.layout.Iter(extent: Expr, stride: Expr, axis: Axis | str)

One extent-and-stride factor contributing to a layout axis.

Parameters:

  • extent (Expr) – Number of positions in this factor.

  • stride (Expr) – Increment on the output axis for each position.

  • axis (Axis or str) – Output coordinate receiving the contribution, resolved by name for strings.

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

Map logical indices to memory and ownership coordinates using strided factors.

shard holds the factors for distinct logical elements; replica holds factors for copies of those elements. offset gives an origin per axis. Use from_iters() to construct these fields directly.

Parameters:

spec (layout specification) – Compose S[shape:stride] shard factors, optional R[shape:stride] replica factors, and offsets with +. Integer strides use the memory axis m; stride @ Axis.name selects another axis. Omitted strides use contiguous row-major memory order.

Examples

from tvm_tirx.script import tirx as T

layout = T.TileLayout(T.S[4, 8])  # Contiguous 4-by-8 tile.
assert int(layout.apply(1, 2)["m"]) == 10
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_contiguous(analyzer=None) → bool

Whether memory-only tile strides are contiguous, regardless of origin.

Replicas and non-memory axes are rejected. An optional analyzer supplies existing constraints; symbolic proof is otherwise initialized lazily.

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]]

get_scope() → tuple[ThreadScope, ThreadScope] | None

Get the scope pair of the 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.tirx.layout.ComposeLayout(per_element: int, swizzle_len: int, atom_len: int, tile_layout: TileLayout, swizzle_inner: bool = True)

Apply an XOR swizzle to the memory coordinate of a tile layout.

All three bit counts must be nonnegative.

Parameters:

  • per_element (int) – Number of low address bits left unchanged by the swizzle.

  • swizzle_len (int) – Number of bits in each XOR group; zero disables swizzling.

  • atom_len (int) – Separation between the bit groups; must be at least swizzle_len.

  • tile_layout (TileLayout) – The underlying mapping to memory coordinates.

  • swizzle_inner (bool, optional) – If true, XOR the higher group into the lower group; otherwise XOR the lower group into the higher group. Defaults to true.

Examples

from tvm_tirx.script import tirx as T

tile = T.TileLayout(T.S[8, 8])
layout = T.ComposeLayout(0, 2, 3, tile)
# Memory offset 10 becomes 10 ^ ((10 >> 3) & 3) = 11.
assert int(layout.apply(1, 2)["m"]) == 11

S[...] and R[...] build shard and replica layout specifications, respectively. Named axes such as laneid, warpid, tid_in_wg, TLane, and TCol belong to CUDA. Use T.cuda.Axis.laneid or import axes from tvm_tirx.backend.cuda.axis. Trainium’s P, F, and Bank belong to T.trn.Axis. The generic memory axis remains T.Axis.m.

CUDA memory/instruction layout factories are defined in tvm_tirx.backend.cuda.layout and exposed as T.cuda.layout. Trainium layouts use explicit partition, free and bank axes. For example:

from tvm_tirx.backend.trn.axis import P, F, Bank
from tvm_tirx.tirx.layout import TileLayout, S

# A (128, 512) SBUF tile: partition p, free element f.
sbuf_layout = TileLayout(S[(128, 512) : (1 @ P, 1 @ F)])
# Two (128, 512) PSUM tiles, one per bank.
psum_layout = TileLayout(S[(2, 128, 512) : (1 @ Bank, 1 @ P, 1 @ F)])

Annotation strings are not accepted by generic tensor constructors.

Omitting layout or passing layout=None leaves every storage scope without a layout. layout="default" explicitly requests a compact TileLayout. Ordinary views preserve shape, strides and element offsets without creating a layout. TensorType.effective_layout() returns an explicit layout or infers a temporary compact layout for a global tensor. Tile instructions query it through tensor.ty.effective_layout(), using the full tensor shape before slicing an operand region. The query does not change the tensor type or its layout=None field. Noncompact global tensors and other storage scopes require an explicit layout when the instruction uses one.

Represent tensor layouts, iteration axes, and mappings between logical and physical storage.