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_majorcomptime 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 computed from the shape.
Parameters
- shape_types (
TypeList): The types for the shape dimensions.
RowMajorLayout
comptime RowMajorLayout[*shape_types: CoordLike] = Layout[shape_types, *?]
A Layout with row-major (C-order) strides computed from the shape.
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]] = False if (TypeList[L.__shape_types].size != TypeList[values].size) else #kgen.param_list.reduce(values, base=True, reducer=[PrevV: Bool, VA: KGENParamList[CoordLike], idx: __mlir_type.index] PrevV and True if not values[idx].is_tuple.__bool__() else False if not L.__shape_types[idx].is_tuple.__bool__() else False if (TypeList[L.__shape_types[idx].ParamListType].size != TypeList[values[idx].ParamListType].size) else #kgen.param_list.reduce(values[idx].ParamListType, base=True, reducer=[PrevV: Bool, VA: KGENParamList[CoordLike], idx: __mlir_type.index] PrevV and True if not values[idx].ParamListType[idx].is_tuple.__bool__() else False if not L.__shape_types[idx].ParamListType[idx].is_tuple.__bool__() else False if (TypeList[L.__shape_types[idx].ParamListType[idx].ParamListType].size != TypeList[values[idx].ParamListType[idx].ParamListType].size) else #kgen.param_list.reduce(values[idx].ParamListType[idx].ParamListType, base=True, reducer=[PrevV: Bool, VA: KGENParamList[CoordLike], idx: __mlir_type.index] PrevV and True if not values[idx].ParamListType[idx].ParamListType[idx].is_tuple.__bool__() else False if not L.__shape_types[idx].ParamListType[idx].ParamListType[idx].is_tuple.__bool__() else (TypeList[L.__shape_types[idx].ParamListType[idx].ParamListType[idx].ParamListType].size == TypeList[values[idx].ParamListType[idx].ParamListType[idx].ParamListType].size))))
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): 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 = tileouter_shape = shape / tileinner_stride = original strideouter_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): 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. -
row_major: Creates a row-major layout from a shapeCoord. -
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.
Was this page helpful?
Thank you! We'll create more content like this.
Thank you for helping us improve!