Mojo struct
TmemArrayType
@register_passable(trivial)
struct TmemArrayType[dtype: DType, layout: Layout, num_tiles: Int, *, cta_group: Int = 1]
Array of tiles in Tensor Memory (TMEM).
Similar to SMemArray but for TMEM-resident tiles. Provides indexed access to a contiguous array of TmemTensor tiles.
Compile-time constants: Tile: TmemTensor type for each tile. tile_stride: Columns per tile (derived from layout.size()). num_cols: Total TMEM columns used (num_tiles × tile_stride).
Parameters
- dtype (
DType): Element dtype for tiles. - layout (
Layout): Layout of each tile. - num_tiles (
Int): Number of tiles in the array. - cta_group (
Int): CTA group size (1 or 2).
Fields
- base_addr (
Int):
Implemented traits
AnyType,
Copyable,
ImplicitlyCopyable,
ImplicitlyDestructible,
Movable,
RegisterType,
TrivialRegisterType
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
num_cols
comptime num_cols = (num_tiles * TmemArrayType[dtype, layout, num_tiles, cta_group=cta_group].tile_stride)
Tile
comptime Tile = TmemTensor[dtype, layout, cta_group=cta_group]
tile_stride
comptime tile_stride = layout.shape[1].value()
Methods
__init__
__init__(base_addr: Int) -> Self
Initialize array at the given TMEM base address.
__getitem__
__getitem__[T: Intable](self, index: T) -> TmemArrayType[dtype, layout, num_tiles, cta_group=cta_group].Tile
Get tile at the given index.
Returns:
TmemArrayType
Was this page helpful?
Thank you! We'll create more content like this.
Thank you for helping us improve!