For the complete documentation index, see llms.txt. Markdown versions of all pages are available by appending .md to any URL (e.g. /max/get-started.md).
Mojo struct
TileLoader
struct TileLoader[tma_origin: ImmOrigin, dtype: DType, tile_layout: TensorLayout, desc_layout: TensorLayout, /, *, cta_group: Int]
TMA tile loader parameterized on new Layout types.
Uses TmaOpType to derive the TMATensorTile type from new Layout. Accepts TileTensor destinations.
Parameters
- tma_origin (
ImmOrigin): Origin of the TMA descriptor pointer. - dtype (
DType): Element data type. - tile_layout (
TensorLayout): Layout of the tile loaded into shared memory. - desc_layout (
TensorLayout): Layout of the TMA descriptor. - cta_group (
Int): CTA group size (1 or 2 for SM100 2-SM MMA).
Fields
- tma_op (
TileLoader[tma_origin, dtype, tile_layout, desc_layout, cta_group=cta_group].TmaOpPtr): - multicast_mask (
UInt16):
Implemented traits
AnyType,
Copyable,
ImplicitlyCopyable,
ImplicitlyDeletable,
Movable,
RegisterPassable,
TrivialRegisterPassable
comptime members
TmaOp
comptime TmaOp = TMATensorTile[dtype, tile_layout.rank, _to_index_list[tile_layout](), _to_index_list[tile_layout.rank, desc_layout]()]
TmaOpPtr
comptime TmaOpPtr = Pointer[TMATensorTile[dtype, tile_layout.rank, _to_index_list[tile_layout](), _to_index_list[tile_layout.rank, desc_layout]()], tma_origin, _safe=True]
Methods
__init__
def __init__[tma_op_type: AnyType](tma_op: Pointer[tma_op_type, tma_origin, _safe=True], multicast_mask: UInt16) -> Self
Accepts any TMA pointer. Rebinds to the loader's derived type.
Parameters:
- tma_op_type (
AnyType): Compile-time type of the passed TMA descriptor pointer.
Args:
- tma_op (
Pointer[tma_op_type, tma_origin, _safe=True]): Pointer to the TMA descriptor. - multicast_mask (
UInt16): Multicast mask for cluster distribution.
load
def load[LayoutType: TensorLayout](self, dest: TileTensor[dtype, LayoutType, MutAnyOrigin, address_space=AddressSpace.SHARED], ref[AddressSpace._value] barrier: SharedMemBarrier, k_coord: Int, row_coord: Int)
Load a tile using TMA async multicast load.
Parameters:
- LayoutType (
TensorLayout): Layout type of the destination TileTensor.
Args:
- dest (
TileTensor[dtype, LayoutType, MutAnyOrigin, address_space=AddressSpace.SHARED]): Destination SMEM TileTensor tile. - barrier (
SharedMemBarrier): Memory barrier for TMA completion signaling. - k_coord (
Int): K dimension coordinate in global memory (elements). - row_coord (
Int): Row coordinate (M for A, N for B) in global memory (elements).
Was this page helpful?
Thank you! We'll create more content like this.
Thank you for helping us improve!