tilus.ir.RegisterLayout

class tilus.ir.RegisterLayout[source]

Layout for register tensor.

shape

The shape of the layout, which is the shape of the register tensor.

Type:

tuple[int, …]

mode_shape

The size of each mode.

Type:

tuple[int, …]

spatial_modes

The spatial modes.

Type:

tuple[int, …]

local_modes

The local modes.

Type:

tuple[int, …]

Register Layout Operations

auto_local_spatial(num_threads, shape)

Create a local(...).spatial(...) layout

column_local(*shape)

Create a local layout in column-major order.

column_spatial(*shape)

Create a spatial layout in column-major order.

compose(outer, inner)

Compose two layouts together.

concat(lhs, rhs)

Concatenate two layouts.

divide(lhs, rhs)

Divide two layouts.

flatten(layout[, start_dim, end_dim])

Flatten the layout over the dimensions between start_dim and end_dim, inclusive.

local(*shape[, ranks])

Create a local layout.

permute(layout, dims)

Permute the dimensions of the layout.

reduce(layout, dims, *[, keepdims])

Reduce the layout over the given dimensions.

register_layout(shape, mode_shape, ...)

Create a register layout with the given shape, mode shape, spatial modes, and local modes.

reshape(layout, shape)

Reshape the layout to the given shape.

spatial(*shape[, ranks])

Create a spatial layout.

squeeze(layout, dims)

Squeeze the layout over the given dimensions.

unsqueeze(layout, dims)

Unsqueeze the layout over the given dimensions.