For the complete documentation index, see llms.txt. Markdown versions of all pages are available by appending .md to any URL (e.g. /get-started.md).

Mojo module

tile_layout

Provides a mixed compile-time/runtime layout system for tensor memory mapping.

This module provides a layout system where some dimensions can be known at compile time and others determined at runtime, enabling ergonomic layout definitions while maintaining performance through compile-time specialization.

Key components:

• TensorLayout: Trait defining the interface for all mixed layouts.
• Layout: Primary struct implementing a layout with mixed compile-time and runtime dimensions.
• row_major: Create a row-major layout from a shape.
• col_major: Create a column-major layout from a shape.
• blocked_product: Create a hierarchical blocked layout from block and tiler layouts.
• zipped_divide: Divide a layout into inner and outer components by a tile shape.
• coalesce: Simplify a layout by merging dimensions with contiguous strides.

You can import these APIs from the layout package:

from layout.tile_layout import Layout, TensorLayout, row_major, col_major

comptime values

BlockedProductLayout

comptime BlockedProductLayout[BlockLayoutType: TensorLayout, TilerLayoutType: TensorLayout, coalesce_output: Bool = False] = Layout[*?, *?]

Type alias for blocked product layout.

Creates a hierarchical layout by combining a block (inner) layout with a tiler (outer) layout. The result zips corresponding dimensions so that each mode i pairs block[i] with tiler[i]:

• shape[i] = (block.shape[i], tiler.shape[i])
• stride[i] = (block.stride[i], block.cosize * tiler.stride[i])

When coalesce_output is True, contiguous inner/outer pairs per mode are merged into flat dimensions (block_shape[i] * block_stride[i] == outer_stride[i]). This corresponds to the old blocked_product(..., coalesce_output=True) with keep_rank=True.

For fully-static layouts, this can be used directly at the type level:

comptime result = BlockedProductLayout[type_of(block), type_of(tiler)]()
comptime coalesced = BlockedProductLayout[
    type_of(block), type_of(tiler), coalesce_output=True
]()

Parameters

• ​BlockLayoutType (TensorLayout): The inner block layout type.
• ​TilerLayoutType (TensorLayout): The outer tiler layout type.
• ​coalesce_output (Bool): Whether to coalesce contiguous modes. Default is False.

CoalesceLayout

comptime CoalesceLayout[LayoutType: TensorLayout] = Layout[*?, *?]

Type alias for the result of coalesce.

Simplifies a layout by merging dimensions with contiguous strides. Adjacent flattened dimensions where prev_shape * prev_stride == current_stride are combined into a single dimension. Shape-1 dimensions are dropped.

For fully-static layouts, this can be used directly at the type level:

comptime result = CoalesceLayout[type_of(my_layout)]()

Parameters

• ​LayoutType (TensorLayout): The input layout type (must have all dimensions known at compile time).

ColMajorLayout

comptime ColMajorLayout[shape_types: TypeList[values]] = Layout[shape_types, *?]

A Layout with column-major (Fortran-order) strides for a flat shape. For nested shapes use ColMajorNestedLayout.

Parameters

• ​shape_types (TypeList[values]): The types for the shape dimensions.

ColMajorNestedLayout

comptime ColMajorNestedLayout[shape_types: TypeList[values]] = Layout[shape_types, *?]

A Layout with column-major strides for a nested shape (CuTe semantics).

For shape ((a, b), (c, d)) the strides are ((1, a), (a*b, a*b*c)): col-major over the flattened shape, re-nested.

Parameters

• ​shape_types (TypeList[values]): The types for the shape dimensions.

RowMajorLayout

comptime RowMajorLayout[*shape_types: CoordLike] = Layout[shape_types, *?]

A Layout with row-major (C-order) strides for a flat shape. For nested shapes use RowMajorNestedLayout.

Parameters

• ​*shape_types (CoordLike): The types for the shape dimensions.

RowMajorNestedLayout

comptime RowMajorNestedLayout[*shape_types: CoordLike] = Layout[shape_types, *?]

A Layout with row-major strides for a nested shape (CuTe semantics).

For shape ((a, b), (c, d)) the strides are ((b*c*d, c*d), (d, 1)): row-major over the flattened shape, re-nested.

Parameters

• ​*shape_types (CoordLike): The types for the shape dimensions.

UpcastLayout

comptime UpcastLayout[LayoutType: TensorLayout, factor: Int] = Layout[*?, *?]

Type alias for the result of upcast.

Fuses factor consecutive elements per dimension, producing a layout with coarser granularity. For fully-static layouts, this can be used directly at the type level without calling upcast:

comptime result = UpcastLayout[type_of(my_layout), 2]()

Parameters

• ​LayoutType (TensorLayout): The input layout type.
• ​factor (Int): The number of consecutive elements to fuse.

WeaklyCompatible

comptime WeaklyCompatible[L: TensorLayout, C: TypeList[values]] = ParameterList.all_satisfies[comptime[Elt: Bool] Elt]() if (Int(len(L.__shape_types)) == Int(len(values))) else (Int(len(L.__shape_types)) == Int(len(values)))

Check structural compatibility between a layout's shape and coordinate types.

Returns True if compatible, False otherwise. A scalar coordinate element is always compatible. A tuple coordinate element requires the corresponding layout shape element to also be a tuple of the same length, with all sub-elements recursively compatible. Handles up to 4 levels of nesting; beyond that, only length equality is verified.

Parameters

• ​L (TensorLayout): The layout type whose shape structure is checked.
• ​C (TypeList[values]): The coordinate element types to check against.

ZippedDivideLayout

comptime ZippedDivideLayout[LayoutType: TensorLayout, tile: TypeList[values]] = Layout[*?, *?]

Type alias for the result of zipped_divide.

Splits a layout into inner (tile-sized) and outer (number-of-tiles) components. The result is a 2-level hierarchical layout where:

• inner_shape = tile
• outer_shape = shape / tile
• inner_stride = original stride
• outer_stride = stride * tile

For fully-static layouts, this can be used directly at the type level:

comptime result = ZippedDivideLayout[type_of(my_layout), tile.element_types]()

Parameters

• ​LayoutType (TensorLayout): The input layout type.
• ​tile (TypeList[values]): Shape types of the tile used to divide the layout.

Structs

• ​Layout: A layout that supports mixed compile-time and runtime dimensions.

Traits

• ​TensorLayout: Trait defining the interface for mixed compile-time/runtime layouts.

Functions

• ​blocked_product: Creates a blocked layout by combining a block and tiler layout.
• ​coalesce: Simplifies a layout by merging contiguous dimensions.
• ​col_major: Create a column-major layout from variadic arguments.
• ​col_major_nested: Creates a column-major layout from a nested shape Coord.
• ​row_major: Creates a row-major layout from a shape Coord.
• ​row_major_nested: Creates a row-major layout from a nested shape Coord.
• ​upcast: Fuses consecutive elements in a layout to create a coarser layout.
• ​zipped_divide: Divides a layout into inner (tile) and outer (number-of-tiles) parts.