Mojo module

layout

Provides a high-performance tensor layout system for memory mapping and indexing.

The layout module implements a comprehensive system for describing memory layouts of multi-dimensional tensors, enabling efficient mapping between logical tensor coordinates and physical memory locations. This is a critical component for high-performance tensor operations in machine learning and scientific computing. These low-level primitives require careful use to avoid errors. Understanding the relationship between tensor shapes, strides, and memory layout is essential for effective use.

Key components:

LayoutTrait: Core trait defining the interface for all layout types
Layout: Primary struct implementing memory layout with shape and stride information
Layout algebra: Functions for composing, dividing, and transforming layouts
Tiling operations: Functions for hierarchical decomposition of layouts

Performance features:

Zero-cost abstractions for mapping between logical and physical indices
Support for both compile-time and runtime-determined shapes
Efficient memory access patterns through layout transformations
Hierarchical tiling for cache-friendly memory access

Common use cases:

Defining memory layouts for tensors with different storage formats (row-major, column-major)
Implementing efficient tensor operations with optimal memory access patterns
Supporting hardware-specific memory layouts for accelerators
Enabling zero-copy tensor views and reshaping operations

Example:

from layout import Layout, IntTuple
from layout.layout import blocked_product

# Create a 3x4 row-major layout
var layout = Layout.row_major(3, 4)

# Access the memory location for logical coordinates (1, 2)
var memory_idx = layout(IntTuple(1, 2))

# Create a tiled layout for blocked matrix multiplication
var tiled = blocked_product(layout, Layout(IntTuple(2, 2)))
from layout import Layout, IntTuple
from layout.layout import blocked_product

# Create a 3x4 row-major layout
var layout = Layout.row_major(3, 4)

# Access the memory location for logical coordinates (1, 2)
var memory_idx = layout(IntTuple(1, 2))

# Create a tiled layout for blocked matrix multiplication
var tiled = blocked_product(layout, Layout(IntTuple(2, 2)))

Aliases

LayoutList = List[Layout]:

Structs

Layout: Represents a memory layout for multi-dimensional data.

Traits

LayoutTrait: Defines the interface for mapping between logical coordinates and memory indices.

Functions

apply_tiler: Applies a layout transformation function to each element of a layout with a tiler.
blocked_product: Creates a blocked layout by combining two layouts.
coalesce: Simplifies a layout by combining dimensions with contiguous strides.
complement: Computes the complement layout for a given layout.
composition: Composes two layouts to create a new layout.
cosize: Returns the size of the memory region spanned by the layout.
downcast: Splits elements in a layout to create a finer layout without changing the total number of elements so that the alignment is preserved.
expand_modes_alike: Aligns two shape-stride pairs to have the same hierarchical structure.
expand_strides: Expands a scalar stride into a stride tuple matching a shape tuple.
format_layout: Formats a 2D layout as a table and writes it to the specified writer.
hierarchical_unzip: Hierarchically unzips a layout according to a list of layouts.
is_row_major: Checks if a layout has row-major ordering for the specified rank.
logical_divide: Divides a layout into blocks according to another layout.
logical_product: Creates a product of two layouts.
make_layout: Creates a composite layout by concatenating multiple layouts.
make_ordered_layout: Creates a layout with strides ordered according to a specified traversal order.
MakeLayoutList: Creates a list containing two layouts.
MakeTileLayoutList: Creates a list of layouts for tiling operations.
print_layout: Prints a 2D layout to the standard output.
right_inverse: Creates a right inverse of a layout.
size: Returns the total number of elements in the layout's domain.
sublayout: Creates a sublayout by selecting specific dimensions from a layout.
tile_to_shape: Creates a layout by tiling a base layout to match a target shape.
upcast: Fuses consecutive elements in a layout to create a coarser layout.
zip_modes: Combines corresponding modes from two layouts.
zipped_divide: Divides a layout into blocks according to another layout.

Aliases​

Structs​

Traits​

Functions​

Aliases

Structs

Traits

Functions