For the complete Mojo documentation index, see llms.txt. Markdown versions of all pages are available by appending .md to any URL (e.g. /docs/manual/basics.md).
Introduction to TileTensor
Mojo's layout package provides a number of APIs for
working with dense multidimensional arrays, which simplify writing algorithms
for handling linear algebra.
This section discusses the TileTensor abstraction and its related types:
TileTensoris a flexible tensor type that combines atile_layout.Layoutand a pointer to data.- The
tile_layout.Layoutstruct describes an arrangement of data in memory. A layout is a function that maps a set of logical coordinates (like (x, y) in a two-dimensional array) to a linear index value (such as an offset into a memory buffer). Layouts can describe simple row-major or column-major matrices, but also more complex organizations, like a matrix subdivided into tiles. Coordis a tuple-like container for storing integer coordinates that supports both compile-time and run-time values. It's used for defining logical coordinates and layout shapes, among other things.
The layout package also includes the older LayoutTensor type, and its
associated IntTuple-based layout type, layout.Layout. TileTensor is
essentially a new version of LayoutTensor that addresses several of the
shortcomings of that type. Most of the concepts are the same for both the old
and new types. However, TileTensor doesn't yet implement all the features
supported by LayoutTensor. Here are some guidelines for when to use each type:
-
Use
TileTensorwhen you need layouts with run-time dimensions, nested layouts, or when working with the newerCoord-based layout system. TheCoord-based layout system minimizes the memory footprint of the layout struct, which can help reduce register pressure on GPUs. -
Use
LayoutTensorwhen you need established operations likecopy_dma(),collective_load(), or compatibility with existing code usingIntTuple-based layouts.
You can convert a TileTensor into a LayoutTensor, so you can
use both types if necessary. Likewise, you can easily convert a
Coord-based layout into an IntTuple-based layout.
For more information on LayoutTensor, see the pages on
IntTuple-based layouts and
LayoutTensor.