tvm_tirx.tirx.TileLayout

class tvm_tirx.tirx.TileLayout(spec)

Bases: Layout

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

property shard

replicate: list[Iter]

exclude: list[tuple[Axis, Expr]]

static from_iters(shard=(), replica=(), offset=None)

Construct a TileLayout from pre-built Iter objects.

Parameters:

• shard (Sequence[Iter])

• replica (Sequence[Iter])

• offset (dict[Axis | str, Expr] | None)

Return type:

TileLayout

is_contiguous(analyzer=None)

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.

Return type:

bool

is_trivial()

Check if the layout is trivial.

Return type:

bool

group(shape)

Group the current layout by the given shape.

Parameters:

shape (List[Expr]) – The shape to group by

Returns:

Tuple[Layout, List[int]] – The grouped layout and the separators

Return type:

tuple[Layout, list[int]]

get_scope()

Get the scope pair of the layout.

Return type:

tuple[ThreadScope, ThreadScope] | None

permute_dims(perm)

Permute the dimensions of the layout.

Parameters:

perm (list[int])

Return type:

TileLayout

permute_by_groups(seps, perm)

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.

Return type:

TileLayout

property offset

property replica