tilelang.layout.layout

Wrapping Layouts.

Classes

Layout

Module Contents

class tilelang.layout.layout.Layout(shape, forward_fn)

Bases: tvm.ir.Node

property index

Property to retrieve the forward index of the layout.

Returns:

The computed forward index expression(s).

Return type:

PrimExpr or List[PrimExpr]

get_input_shape()

Get the input shape of the layout.

Returns:

The shape of the input layout.

Return type:

List[int]

get_output_shape()

Get the output shape of the layout.

Returns:

The shape of the output layout.

Return type:

List[int]

get_forward_vars()

Retrieve the iteration variables associated with the layout.

Returns:

A list of iteration variables that define the layout transformation.

Return type:

List[IterVar]

get_forward_index()
map_forward_index(indices)

Compute the forward index mapping for a given set of input indices.

Parameters:

indices (list of PrimExpr) – The input indices to be mapped to their corresponding output indices.

Returns:

The mapped index expression for the provided input indices.

Return type:

PrimExpr

repeat(dim, factor)

Repeat a layout along a single input dimension.

This is useful for building a larger layout by tiling an “atom” layout. Conceptually, repeating on dimension dim with factor constructs a new layout L' such that:

L'(*idx) = [idx[dim] // extent_dim] + L(idx with idx[dim] % extent_dim)

where extent_dim is the original extent of the repeated dimension.

Parameters:

  • dim (int) – The input dimension to repeat (0-based, supports negative indexing).

  • factor (int) – The repeat factor. Must be a positive integer.

Returns:

A new Layout with the repeated input shape and an extra leading output dimension representing the repeat-group index.

Return type:

Layout

expand(leading_shape)

Expand (lift) this layout by prepending new leading input dimensions.

The new leading dimensions are forwarded unchanged to the output, and the original layout is applied to the remaining trailing dimensions.

Example

Given a 2D layout L over [J, K], you can lift it to a 3D layout over [I, J, K] by:

L3 = L.expand([I])
# [i, j, k] -> [i, *L(j, k)]
Parameters:

leading_shape (int or Sequence[int or PrimExpr]) – The shape of the new leading dimensions to prepend. Use an empty list/tuple for a no-op.

Returns:

A new Layout with input shape leading_shape + input_shape and output indices [leading_dims] + old_forward_index.

Return type:

Layout

inverse()

Compute the inverse of the current layout transformation.

Returns:

A new Layout object representing the inverse transformation.

Return type:

Layout

reshape(shape, rescale_num=1, rescale_den=1)

Reshape the input shape of the layout.

Parameters:

  • shape (list[PrimExpr] or list[int]) – The new input shape.

  • rescale_num (int) – Rescale numerator for element size changes.

  • rescale_den (int) – Rescale denominator for element size changes.

Return type:

Layout

is_equal(other)

Check if the current layout is equal to another layout.

Parameters:

other (Layout) – The layout to compare with.

Return type:

bool

__call__(*args)
Parameters:

args (list[tvm.tir.PrimExpr])

Return type:

tvm.tir.PrimExpr

__repr__()