tvm_tirx.tirx.Layout

tvm_tirx.tirx.Layout#

class tvm_tirx.tirx.Layout(*args, **kwargs)#

Bases: Object

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.

Parameters:

  • args (Any)

  • kwargs (Any)

check_well_formed()#

Check whether the layout is well formed.

Raises:

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

Return type:

None

size(axis_name=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=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, shape=None, analyzer=None)#

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

analyzer optionally supplies caller-owned analysis state for repeated queries. Its premises may simplify the result; this method adds no persistent premises and does not retain the analyzer. Concurrent calls must use separate analyzers. Without it, a temporary analyzer with the DEFAULT-strength Z3 fallback disabled is created. Pass an analyzer explicitly to control its proof settings.

Returns:

Dict[str, Expr] – The mapped output (axis name -> value on the axis)

Parameters:
Return type:

dict[str, Expr]

canonicalize()#

Canonicalize the layout by simplifying and fusing iterators where possible.

Returns:

Layout – The canonicalized layout

Return type:

Layout

tile(outer, outer_shape, inner_shape)#

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:

Union[TileLayout, ComposeLayout] – The resulting tiled layout

Return type:

TileLayout | ComposeLayout

tile_to(to_shape, current_shape)#

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

Return type:

Layout

is_trivial()#

Check if the layout is trivial.

Return type:

bool

unpack(num, axis='m')#

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:

Layout – The unpacked layout

Return type:

Layout

pack(num, axis='m')#

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:

Layout – The packed layout

Return type:

Layout