tilus.ir.SharedLayout

tilus.ir.SharedLayout

class tilus.ir.SharedLayout[source]

The layout for shared tensor.

We use three components to describe a shared tensor layout: the shape, the mode shape, and the mode strides.

The mode shape and mode strides are used to describe how to split each dimension into multiple sub-dimensions (modes), and the strides of each mode.

For example, consider a shape of (64, 32), we can split the first dimension into two sub-dimensions (modes) of size 8 and 8, and the second dimension into two sub-dimensions (modes) of size 16 and 2. The mode shape would be (8, 8, 16, 2). We can have strides for each mode, for example, (256, 2, 16, 1). Then given the indices (i, j), we can compute the indices in the sub-dimensions (i1, i2, j1, j2) where i1 = i // 8, i2 = i % 8, j1 = j // 2, j2 = j % 2. The offset can be computed as: offset = i1 * 256 + i2 * 2 + j1 * 16 + j2 * 1. To get the final offset in the shared tensor, we can use the formula: (i, j) => ((i // 8) * 256) + ((i % 8) * 2) + ((j // 2) * 16) + ((j % 2) * 1).

shape

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

Type:

tuple[int, …]

mode_shape

We can split each dimension into multiple sub-dimensions (modes).

Type:

tuple[int, …]

mode_strides

The strides of each mode.

Type:

tuple[int, …]

swizzle

The swizzle function to apply on the final offset. If None, no swizzling is applied.

Type:

Optional[Swizzle]

__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, mode_shape, mode_strides, optional_swizzle)[source]

Create a SharedLayout from shape, mode_shape, and mode_strides.

Parameters:

  • shape (Sequence[int]) – The shape of the shared tensor.

  • mode_shape (Sequence[int]) – The mode shape of the shared tensor.

  • mode_strides (Sequence[int]) – The mode strides of the shared tensor.

  • swizzle (Optional[Swizzle]) – The swizzle function to apply on the final offset. If None, no swizzling is applied.

  • optional_swizzle (Swizzle | None)

Returns:

ret – The created SharedLayout.

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)

Compose multiple shared layouts together.
Contents