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.

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[block_shape_types: Variadic[CoordLike], block_stride_types: Variadic[CoordLike], tiler_shape_types: Variadic[CoordLike], tiler_stride_types: Variadic[CoordLike]] = Layout[Coord[block_shape_types], Coord[tiler_shape_types], Coord[block_stride_types], Coord[#kgen.variadic.reduce(tiler_stride_types, base=, reducer=[PrevV: Variadic[CoordLike], VA: Variadic[CoordLike], idx: __mlir_type.index] #kgen.variadic.concat(PrevV, ComptimeInt[(VA[idx].static_value * Coord[block_shape_types].static_product)]))]]

Type alias for blocked product layout.

Creates a hierarchical layout by combining a block (inner) layout with a tiler (outer) layout. This is useful for creating tiled memory access patterns.

The result is a layout where:

• inner_shape = block.shape (dimensions within each tile)
• outer_shape = tiler.shape (how many tiles in each dimension)
• inner_stride = block.stride (stride within a tile)
• outer_stride = block.cosize * tiler.stride (stride between tiles)

Parameters

• ​block_shape_types (Variadic): Shape types for the inner block.
• ​block_stride_types (Variadic): Stride types for the inner block.
• ​tiler_shape_types (Variadic): Shape types for the outer tiler (number of blocks).
• ​tiler_stride_types (Variadic): Stride types for the outer tiler.

ColMajorLayout

comptime ColMajorLayout[*shape_types: CoordLike] = Layout[shape_types, #kgen.variadic.reduce(shape_types, base=, reducer=[PrevV: Variadic[CoordLike], VA: Variadic[CoordLike], idx: __mlir_type.index] #kgen.variadic.concat(PrevV, ComptimeInt[1] if (idx == 0)._mlir_value else RuntimeInt[VA[(add idx, -1)].DTYPE if VA[(add idx, -1)].is_static_value.__bool__().__invert__()._mlir_value else PrevV[(add idx, -1)].DTYPE] if VA[(add idx, -1)].is_static_value.__bool__().__invert__() if VA[(add idx, -1)].is_static_value.__bool__().__invert__()._mlir_value else PrevV[(add idx, -1)].is_static_value.__bool__().__invert__() else ComptimeInt[(VA[(add idx, -1)].static_value * PrevV[(add idx, -1)].static_value)]))]

A Layout with column-major (Fortran-order) strides computed from the shape.

Parameters

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

RowMajorLayout

comptime RowMajorLayout[*shape_types: CoordLike] = Layout[shape_types, #kgen.variadic.reduce(#kgen.variadic.reduce(shape_types, base=, reducer=[PrevV: Variadic[CoordLike], VA: Variadic[CoordLike], idx: __mlir_type.index] #kgen.variadic.concat(PrevV, VA[(add (mul idx, -1), len(VA), -1)])), base=, reducer=[PrevV: Variadic[CoordLike], VA: Variadic[CoordLike], idx: __mlir_type.index] #kgen.variadic.concat(ComptimeInt[1] if (idx == 0)._mlir_value else RuntimeInt[VA[(add idx, -1)].DTYPE if VA[(add idx, -1)].is_static_value.__bool__().__invert__()._mlir_value else PrevV[0].DTYPE] if VA[(add idx, -1)].is_static_value.__bool__().__invert__() if VA[(add idx, -1)].is_static_value.__bool__().__invert__()._mlir_value else PrevV[0].is_static_value.__bool__().__invert__() else ComptimeInt[(VA[(add idx, -1)].static_value * PrevV[0].static_value)], PrevV))]

A Layout with row-major (C-order) strides computed from the shape.

Parameters

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

ZippedDivideLayout

comptime ZippedDivideLayout[shape_types: Variadic[CoordLike], stride_types: Variadic[CoordLike], tile: Variadic[CoordLike]] = Layout[Coord[tile], Coord[#kgen.variadic.reduce(shape_types, base=, reducer=[PrevV: Variadic[CoordLike], VA: Variadic[CoordLike], idx: __mlir_type.index] #kgen.variadic.concat(PrevV, ComptimeInt[(VA[idx].static_value // tile[idx].static_value)]))], Coord[stride_types], Coord[#kgen.variadic.reduce(stride_types, base=, reducer=[PrevV: Variadic[CoordLike], VA: Variadic[CoordLike], idx: __mlir_type.index] #kgen.variadic.concat(PrevV, ComptimeInt[(VA[idx].static_value * tile[idx].static_value)]))]]

Type alias for the result of zipped_divide.

Splits a layout into inner (tile-sized) and outer (number-of-tiles) components.

Parameters

• ​shape_types (Variadic): Shape types of the original layout.
• ​stride_types (Variadic): Stride types of the original layout.
• ​tile (Variadic): 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.
• ​col_major: Create a column-major layout from a shape.
• ​row_major: Creates a row-major layout from a shape Coord.
• ​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!

😔 What went wrong?

Some code doesn’t work

It includes inaccurate information

It's missing information I need

It was difficult to understand

Other

Submit