Mojo struct

Layout

struct Layout

Represents a memory layout for multi-dimensional data.

The Layout struct is the primary implementation of the LayoutTrait, providing a concrete representation of memory layouts using shape and stride information. It maps between logical coordinates and linear memory indices, enabling efficient access to multi-dimensional data.

A Layout consists of:

shape: Defines the dimensions of the logical coordinate space
stride: Defines the step sizes in memory for each dimension

The Layout struct supports various operations including:

Creation of row-major and column-major layouts
Conversion between coordinates and indices
Composition with other layouts
Iteration over sub-layouts

Layouts can be hierarchical, with nested shapes and strides, allowing for complex memory access patterns like blocked or tiled layouts.

Fields

shape (IntTuple): The dimensions of the layout. This field defines the size of each dimension in the logical coordinate space. For example, a shape of (3, 4) represents a 3x4 grid of elements.
stride (IntTuple): The memory step sizes for each dimension. This field defines how many elements to skip in memory when moving one unit in each dimension. For example, in a row-major 3x4 layout, the strides might be (4, 1), meaning moving one unit in the first dimension requires skipping 4 elements in memory, while moving one unit in the second dimension requires skipping 1 element.

Implemented traits

AnyType, Copyable, Defaultable, EqualityComparable, ImplicitlyCopyable, Iterable, LayoutTrait, Movable, Sized, Stringable, UnknownDestructibility, Writable

Aliases

`copyinitis_trivial`

alias __copyinit__is_trivial = False

`delis_trivial`

alias __del__is_trivial = False

`moveinitis_trivial`

alias __moveinit__is_trivial = True

`has_shape`

alias has_shape = True

Indicates whether the layout has a valid shape.

`IteratorType`

alias IteratorType[iterable_mut: Bool, //, iterable_origin: Origin[iterable_mut]] = _LayoutIter[origin_of((muttoimm iterable_origin._mlir_origin))]

Parameters

iterable_mut (Bool):
iterable_origin (Origin):

Methods

`init`

__init__(out self)

Initializes an empty layout with no dimensions.

Creates a layout with empty shape and stride tuples, which can be populated later using append operations.

__init__(out self, shape: IntTuple)

Initializes a layout with the given shape and column-major strides.

Creates a layout with the specified shape and automatically calculates column-major strides (where the first dimension varies fastest in memory).

Args:

shape (IntTuple): The dimensions of the layout.

__init__(out self, shape: IntTuple, stride: IntTuple)

Initializes a layout with the given shape and stride.

Creates a layout with explicitly specified shape and stride values. If an empty stride is provided, column-major strides are calculated.

Args:

shape (IntTuple): The dimensions of the layout.
stride (IntTuple): The memory step size for each dimension, or empty for column-major.

`getitem`

__getitem__(self, index: Int) -> Self

Returns a sub-layout for the specified dimension.

Args:

index (Int): The dimension index to extract.

Returns:

Self: A Layout containing the shape and stride for the specified dimension.

`eq`

__eq__(self, other: Self) -> Bool

Checks if this layout is equal to another layout.

Two layouts are considered equal if they have identical shape and stride tuples.

Args:

other (Self): The layout to compare with.

Returns:

Bool: True if the layouts are equal, False otherwise.

`idx2crd`

idx2crd(self, idx: IntTuple) -> IntTuple

Converts a linear index to logical coordinates.

This is the inverse operation of the call method, mapping from a memory index back to the corresponding logical coordinates.

Args:

idx (IntTuple): The linear index to convert.

Returns:

IntTuple: The logical coordinates corresponding to the given index.

`col_major`

static col_major(*dims: Int) -> Self

Creates a column-major layout with the specified dimensions.

In a column-major layout, the first dimension varies fastest in memory, which is the default layout in languages like Fortran and MATLAB.

Example:

from layout import Layout

# Create a 3x4 column-major layout
var layout = Layout.col_major(3, 4)
# Result: Layout with shape (3,4) and stride (1,3)

Args:

*dims (Int): Variable number of dimension sizes.

Returns:

Self: A column-major Layout with the specified dimensions

static col_major(shape: IntTuple) -> Self

Creates a column-major layout with the specified shape.

In a column-major layout, the first dimension varies fastest in memory, which is the default layout in languages like Fortran and MATLAB.

Example:

from layout import Layout
from layout.int_tuple import IntTuple

# Create a 3x4 column-major layout
var layout = Layout.col_major(IntTuple(3, 4))
# Result: Layout with shape (3,4) and stride (1,3)

Args:

shape (IntTuple): An IntTuple specifying the dimensions.

Returns:

Self: A column-major Layout with the specified shape

static col_major[rank: Int](dims: DimList) -> Self

Creates a col-major layout from a DimList with compile-time rank.

This method creates a col-major layout where the first dimension varies fastest in memory. It handles both known and unknown dimensions at compile time, properly calculating strides for each dimension. If any dimension is unknown, subsequent strides will also be marked as unknown.

Example:

from layout import Layout
from layout.layout import DimList

# Create a col-major layout with compile-time rank
var dims = DimList(3, 4)
var layout = Layout.col_major[2](dims)
# Result: Layout with shape (3,4) and stride (1,3)

Parameters:

rank (Int): The compile-time rank (number of dimensions) of the layout.

Args:

dims (DimList): A DimList containing the dimensions of the layout.

Returns:

Self: A col-major Layout with the specified dimensions and computed strides.

static col_major[rank: Int](tuple: IndexList[rank]) -> Self

Creates a col-major layout from a IndexList with compile-time rank.

Example:

from layout import Layout
from utils import IndexList

# Create a row-major layout with compile-time rank
var idx_list = IndexList[2](3, 4)
var layout = Layout.col_major[2](idx_list)
# Result: Layout with shape (3,4) and stride (1,3)

Parameters:

rank (Int): The compile-time rank (number of dimensions) of the layout.

Args:

tuple (IndexList): An IndexList containing the dimensions of the layout.

Returns:

Self: A col-major Layout with the specified dimensions and computed strides.

`row_major`

static row_major(*dims: Int) -> Self

Creates a row-major layout with the specified dimensions.

In a row-major layout, the last dimension varies fastest in memory, which is the default layout in languages like C, C++, and Python.

Example:

from layout import Layout

# Create a 3x4 row-major layout
var layout = Layout.row_major(3, 4)
# Result: Layout with shape (3,4) and stride (4,1)

Args:

*dims (Int): Variable number of dimension sizes.

Returns:

Self: A row-major Layout with the specified dimensions

static row_major[rank: Int](dims: DimList) -> Self

Creates a row-major layout from a DimList with compile-time rank.

This method creates a row-major layout where the last dimension varies fastest in memory. It handles both known and unknown dimensions at compile time, properly calculating strides for each dimension. If any dimension is unknown, subsequent strides will also be marked as unknown.

Example:

from layout import Layout
from layout.layout import DimList

# Create a row-major layout with compile-time rank
var dims = DimList(3, 4)
var layout = Layout.row_major[2](dims)
# Result: Layout with shape (3,4) and stride (4,1)

Parameters:

rank (Int): The compile-time rank (number of dimensions) of the layout.

Args:

dims (DimList): A DimList containing the dimensions of the layout.

Returns:

Self: A row-major Layout with the specified dimensions and computed strides.

static row_major[rank: Int](tuple: IndexList[rank]) -> Self

Creates a row-major layout from a IndexList with compile-time rank.

Example:

from layout import Layout
from utils import IndexList

# Create a row-major layout with compile-time rank
var idx_list = IndexList[2](3, 4)
var layout = Layout.row_major[2](idx_list)
# Result: Layout with shape (3,4) and stride (4,1)

Parameters:

rank (Int): The compile-time rank (number of dimensions) of the layout.

Args:

tuple (IndexList): An IndexList containing the dimensions of the layout.

Returns:

Self: A row-major Layout with the specified dimensions and computed strides.

static row_major[rank: Int]() -> Self

Creates a row-major layout with unknown values for each axis from a compile-time rank.

Example:

from layout import Layout

var layout = Layout.row_major[2]()
# Result: Layout with shape (UNKNOWN_VALUE, UNKNOWN_VALUE)

Parameters:

rank (Int): The compile-time rank (number of dimensions) of the layout.

Returns:

Self: A row-major Layout with the given rank.

static row_major(shape: IntTuple) -> Self

Creates a row-major layout from an IntTuple of dimensions.

In a row-major layout, the last dimension varies fastest in memory. This method computes the appropriate strides for a row-major layout given the input shape.

Example:

from layout import Layout
from layout.int_tuple import IntTuple

# Create a row-major layout from a shape tuple
var shape = IntTuple(3, 4)
var layout = Layout.row_major(shape)
# Result: Layout with shape (3,4) and stride (4,1)

Args:

shape (IntTuple): An IntTuple containing the dimensions of the layout.

Returns:

Self: A row-major Layout with the specified shape and computed strides.

`make_shape_unknown`

make_shape_unknown[axis: Int = -1](self) -> Self

Creates a new Layout with unknown shape dimensions.

This method creates a copy of the current Layout but marks either all dimensions or a specific dimension as unknown, while preserving the original strides. This is useful for tiling tensors with runtime sizes where the tile's shape is unknown but the memory layout (strides) remains constant.

Example:

from layout import Layout
from layout.int_tuple import IntTuple

# Mark all dimensions as unknown
var layout = Layout(IntTuple(2, 3))
var unknown = layout.make_shape_unknown()
# Result: Layout with shape (?, ?) and original strides

# Mark only first dimension as unknown
var partial = layout.make_shape_unknown[0]()
# Result: Layout with shape (?, 3) and original strides

Parameters:

axis (Int): The dimension to mark as unknown. If UNKNOWN_VALUE (default), all dimensions are marked as unknown.

Returns:

Self: A new Layout with the specified dimension(s) marked as unknown and original strides preserved.

`str`

__str__(self) -> String

Converts the layout to a string representation.

Returns:

String: A string representation of the layout in the format "(shape:stride)".

`write_to`

write_to(self, mut writer: T)

Writes the layout to the specified writer.

Formats the layout as "(shape:stride)" and writes it to the provided writer.

Args:

writer (T): The writer to output the layout representation to.

`len`

__len__(self) -> Int

Returns the number of dimensions in the layout.

Returns:

Int: The number of elements in the shape tuple.

`iter`

__iter__(ref self) -> _LayoutIter[origin_of((muttoimm self_is_origin))]

Returns an iterator over the layout's dimensions.

Each iteration yields a Layout containing the shape and stride for one dimension.

Returns:

_LayoutIter: An iterator over the layout's dimensions.

`size`

size(self) -> Int

Returns the total number of elements in the layout's domain.

Calculates the product of all dimensions in the shape.

Returns:

Int: The total number of elements in the layout.

`cosize`

cosize(self) -> Int

Returns the size of the memory region spanned by the layout.

Calculates the maximum memory index plus one, representing the total memory footprint required by the layout.

Returns:

Int: The size of the memory region required by the layout.

`rank`

rank(self) -> Int

Returns the number of dimensions in the layout.

This is equivalent to len and returns the number of elements in the shape tuple.

Returns:

Int: The number of dimensions in the layout.

`call`

__call__(self, idx: IntTuple) -> Int

Maps logical coordinates to a linear memory index.

This is the core functionality of a layout, converting multi-dimensional coordinates to a linear memory location.

Args:

idx (IntTuple): The logical coordinates to map.

Returns:

Int: The linear memory index corresponding to the given coordinates.

`append`

append(mut self, item: Self)

Appends another layout to this layout.

This method adds the shape and stride from the provided layout to this layout, effectively increasing its dimensionality.

Args:

item (Self): The layout to append to this layout.

`all_dims_known`

all_dims_known(self) -> Bool

Checks if all dimensions in the layout have known values.

A dimension is considered unknown if its shape or stride is set to the special UNKNOWN_VALUE constant.

Returns:

Bool: True if all dimensions have known shape and stride values, False otherwise.

`known_shape`

known_shape(self) -> Bool

Checks if all shape dimensions in the layout have known values.

A dimension is considered unknown if its shape is set to the special UNKNOWN_VALUE constant. This method only checks shapes, not strides.

Returns:

Bool: True if all shape dimensions have known values, False otherwise.

`transpose`

transpose(self) -> Self

Transposes the layout by reversing the order of dimensions.

For an n-dimensional layout, this reverses the order of both shapes and strides. For nested layouts, only the top-level dimensions are transposed, not the hierarchical structure within nested tuples.

Example:

from layout import Layout
from layout.int_tuple import IntTuple

# Simple 2D transpose (row-major to column-major)
var layout = Layout.row_major(3, 4)  # shape (3,4), stride (4,1)
var transposed = layout.transpose()  # shape (4,3), stride (1,4)

# 3D transpose
var layout3d = Layout.row_major(2, 3, 4)  # shape (2,3,4), stride (12,4,1)
var trans3d = layout3d.transpose()        # shape (4,3,2), stride (1,4,12)

# Nested layout - only top level transposed
var nested = Layout(
    IntTuple(IntTuple(2, 3), 4),
    IntTuple(IntTuple(12, 4), 1)
)
var trans_nested = nested.transpose()
# Result: shape (4, (2,3)), stride (1, (12,4))

Returns:

Self: A new Layout with transposed dimensions.

Fields​

Implemented traits​

Aliases​

__copyinit__is_trivial​

__del__is_trivial​

__moveinit__is_trivial​

has_shape​

IteratorType​

Parameters​

Methods​

__init__​

__getitem__​

__eq__​

idx2crd​

col_major​

row_major​

make_shape_unknown​

__str__​

write_to​

__len__​

__iter__​

size​

cosize​

rank​

__call__​

append​

all_dims_known​

known_shape​

transpose​

Fields

Implemented traits

Aliases

`copyinitis_trivial`

`delis_trivial`

`moveinitis_trivial`

`has_shape`

`IteratorType`

Parameters

Methods

`init`

`getitem`

`eq`

`idx2crd`

`col_major`

`row_major`

`make_shape_unknown`

`str`

`write_to`

`len`

`iter`

`size`

`cosize`

`rank`

`call`

`append`

`all_dims_known`

`known_shape`

`transpose`