tilus.ir.GlobalLayout

tilus.ir.GlobalLayout

class tilus.ir.GlobalLayout[source]

The layout for global tensor.

shape

The shape of the global tensor. Each dimension can be an expression of grid-invariant expression, or a constant integer.

Type:

tuple[Expr, …]

size

The storage size of the global 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.

Type:

Expr

axes

The axes of the global 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

An expression that computes the offset of the global tensor based on the axes. Only the axes and grid-invariant variables 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 global tensor with the given indices.

Parameters:

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

Returns:

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

Return type:

Expr

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

Create a global layout with custom mapping.

This method creates a global layout with the given shape, size, and the mapping function for the offset.

Parameters:

  • shape (Sequence[Expr | int]) – The shape of the global tensor. Each dimension can be an expression of grid-invariant expression, or a constant integer.

  • size (Expr) – The storage size of the global 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.

  • f_offset (Callable[[Sequence[Var]], Expr]) – A function that takes a sequence of axes (variables) and returns an expression that computes the offset of the global tensor based on those axes. The axes are variables that represent the index of the corresponding dimension. The function should not use any variables that are not grid-invariant.

Returns:

ret – A new instance of GlobalLayout with the specified shape, size, axes, and offset function.

Return type:

GlobalLayout

Global Layout Creation

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

global_row_major(*shape)

Create a global layout with row-major order.

global_column_major(*shape)

Create a global layout with column-major order.

global_strides(shape, strides)

Create a global layout with specified strides.

global_compose(lhs, rhs, *others)

Compose multiple global layouts.
Contents