tilus.ir.SharedLayout

tilus.ir.SharedLayout

class tilus.ir.SharedLayout[source]

The layout for shared tensor.

shape

The shape of the shared tensor. Each dimension is a constant integer.

Type:

tuple[int, …]

size

The storage size of the shared tensor, in number of elements. If the layout is a compact layout, size should be equal to the product of the shape dimensions. Otherwise, it can be either larger (in case of padding) or smaller (in case of sharing data for different elements) than the product of the shape dimensions. The size must be a constant integer.

Type:

int

axes

The axes of the shared tensor. Each axis is a variable that represents the index of the corresponding dimension. It should have the same length as the shape.

Type:

tuple[Var, …]

offset

The offset expression of the shared tensor based on the axes. It is an expression that computes the offset of the shared tensor based on the axes. Only the axes and variables that are invariant in the lifetime of the given corresponding shared tensor with this layout can be used in the expression.

Type:

Expr

__call__(*indices)[source]

Compute the offset on given indices.

This method computes the offset of an element in the shared tensor with the given indices.

Parameters:

indices (Sequence[Expr]) – The indices of the shared tensor. The length of the indices should match the number of axes in the layout.

Returns:

ret – The computed offset of the shared tensor element at the given indices.

Return type:

Expr

static create(shape, size, f_offset)[source]

Create a shared layout.

This method creates a shared layout with the given shape, size, and a function to compute the offset based on the axes. The shape must be a sequence of constant integers, and the size must be a constant integer that is larger than the maximum possible offset computed by the f_offset function.

Parameters:

  • shape (Sequence[int]) – The shape of the shared tensor. Each dimension is a constant integer.

  • size (int) – The storage size of the shared tensor, in number of elements.

  • f_offset (Callable[[Sequence[Var]], Expr]) – The function that computes the offset of the shared tensor based on the axes. It takes a sequence of axes (variables) and returns an expression that computes the offset. The function must ensure that the size is larger than the maximum possible offset computed by this function.

Returns:

ret – A shared layout with the specified shape, size, axes, and offset.

Return type:

SharedLayout

Shared Layout Creation

We can use the following functions to create a shared layout:

shared_row_major(*shape)

Create a shared layout with row-major order.

shared_column_major(*shape)

Create a shared layout with column-major order.

shared_compose(lhs, rhs, *others)

Compose multiple shared layouts together.
Contents