Mojo struct
TileLoader
@register_passable(trivial)
struct TileLoader[tma_origin: ImmutOrigin, dtype: DType, tile_layout: TensorLayout, desc_layout: TensorLayout, /, *, cta_group: Int]
TMA tile loader parameterized on new Layout types.
Uses TMATile to derive the TMATensorTile type from new Layout. Accepts TileTensor destinations.
Fields
- tma_op (
TileLoader[tma_origin, dtype, tile_layout, desc_layout, cta_group=cta_group].TmaOpPtr): - multicast_mask (
UInt16):
Implemented traits
AnyType,
Copyable,
ImplicitlyCopyable,
ImplicitlyDestructible,
Movable,
RegisterPassable,
TrivialRegisterPassable
comptime members
__copyinit__is_trivial
comptime __copyinit__is_trivial = True
__del__is_trivial
comptime __del__is_trivial = True
__moveinit__is_trivial
comptime __moveinit__is_trivial = True
TmaOp
comptime TmaOp = TMATile[dtype, tile_layout, desc_layout].InnerType
TmaOpPtr
comptime TmaOpPtr = Pointer[TileLoader[tma_origin, dtype, tile_layout, desc_layout, cta_group=cta_group].TmaOp, tma_origin]
Methods
__init__
__init__[tma_op_type: AnyType](tma_op: Pointer[tma_op_type, tma_origin], multicast_mask: UInt16) -> Self
Accepts any TMA pointer. Rebinds to the loader's derived type.
load
load[LayoutType: TensorLayout](self, dest: TileTensor[dtype, LayoutType, MutAnyOrigin, address_space=AddressSpace.SHARED], ref[AddressSpace._value._mlir_value] barrier: SharedMemBarrier, k_coord: Scalar[DType.uint], row_coord: Scalar[DType.uint])
Load a tile using TMA async multicast load.
Was this page helpful?
Thank you! We'll create more content like this.
Thank you for helping us improve!