Chapter 8: 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.