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 3×4 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 3×4 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,
LayoutTrait,
Movable,
Sized,
Stringable,
UnknownDestructibility,
Writable
Aliases
has_shape
alias has_shape = True
Indicates whether the layout has a valid shape.
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.
@implicit
__init__(out self, shape: IntTuple[origin])
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[origin]): The dimensions of the layout.
__init__(out self, shape: IntTuple[origin], stride: IntTuple[origin])
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[origin]): The dimensions of the layout. - stride (
IntTuple[origin]): The memory step size for each dimension, or empty for column-major.
__init__(out self, *, other: Self)
Explicitly constructs a deep copy of the provided layout.
Args:
- other (
Self): The layout to copy.
__getitem__
__getitem__(self, index: Int) -> Self
Returns a sub-layout for the specified dimension.
Args:
- index (
Int): The dimension index to extract.
Returns:
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:
True if the layouts are equal, False otherwise.
__ne__
__ne__(self, other: Self) -> Bool
Checks if this layout is not equal to another layout.
Args:
- other (
Self): The layout to compare with.
Returns:
True if the layouts are not equal, False otherwise.
idx2crd
idx2crd(self, idx: IntTuple[origin]) -> 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[origin]): The linear index to convert.
Returns:
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)
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:
A column-major Layout with the specified dimensions
static col_major(shape: IntTuple[origin]) -> 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)
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[origin]): An IntTuple specifying the dimensions.
Returns:
A column-major Layout with the specified shape
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)
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:
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)
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:
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 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)
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:
- tuple (
IndexList[rank]): An IndexList containing the dimensions of the layout.
Returns:
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)
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:
A row-major Layout with the given rank.
static row_major(shape: IntTuple[origin]) -> 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)
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[origin]): An IntTuple containing the dimensions of the layout.
Returns:
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
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:
A new Layout with the specified dimension(s) marked as unknown and original strides preserved.
copy
copy(self) -> Self
Explicitly constructs a copy of this layout.
Creates a deep copy of the layout, including its shape and stride tuples.
Returns:
A new Layout instance with identical shape and stride values.
__str__
__str__(self) -> String
Converts the layout to a string representation.
Returns:
A string representation of the layout in the format "(shape:stride)".
write_to
write_to[W: Writer](self, mut writer: W)
Writes the layout to the specified writer.
Formats the layout as "(shape:stride)" and writes it to the provided writer.
Parameters:
- W (
Writer): Type parameter representing a Writer implementation.
Args:
- writer (
W): The writer to output the layout representation to.
__len__
__len__(self) -> Int
Returns the number of dimensions in the layout.
Returns:
The number of elements in the shape tuple.
__iter__
__iter__(self) -> _LayoutIter[self]
Returns an iterator over the layout's dimensions.
Each iteration yields a Layout containing the shape and stride for one dimension.
Returns:
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:
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:
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:
The number of dimensions in the layout.
__call__
__call__(self, idx: IntTuple[origin]) -> 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[origin]): The logical coordinates to map.
Returns:
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:
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:
True if all shape dimensions have known values, False otherwise.
Was this page helpful?
Thank you! We'll create more content like this.
Thank you for helping us improve!