Mojo struct
TMemTile
@register_passable(trivial)
struct TMemTile[dtype_: DType, BM: Int, BN: Int]
Fields
- tmem_addr (
UInt32):
Implemented traits
AnyType,
Copyable,
ImplicitlyCopyable,
Movable,
UnknownDestructibility
Aliases
__copyinit__is_trivial
comptime __copyinit__is_trivial = True
__del__is_trivial
comptime __del__is_trivial = True
__moveinit__is_trivial
comptime __moveinit__is_trivial = True
dtype
comptime dtype = dtype_
dtype_size
comptime dtype_size = size_of[TMemTile[dtype_, BM, BN].dtype]()
num_m_tiles
comptime num_m_tiles = (BM // 64)
Methods
__init__
__init__(tmem_addr: UInt32) -> Self
__getitem__
__getitem__(self, i: UInt32) -> Self
offset
allocate_register_tile
static allocate_register_tile[*, num_threads: Int]() -> LayoutTensor[TMemTile[dtype_, BM, BN].dtype, STMatrixLayout[BM, BN, num_threads=num_threads, accum_type_size=TMemTile[dtype_, BM, BN].dtype_size].vec_local_layout, MutAnyOrigin, address_space=AddressSpace.LOCAL, element_layout=STMatrixLayout[BM, BN, num_threads=num_threads, accum_type_size=TMemTile[dtype_, BM, BN].dtype_size].element_layout]
Returns:
store_async
store_async[*, num_threads: Int](self, src: LayoutTensor[TMemTile[dtype_, BM, BN].dtype, STMatrixLayout[BM, BN, num_threads=num_threads, accum_type_size=TMemTile[dtype_, BM, BN].dtype_size].vec_local_layout, MutAnyOrigin, address_space=AddressSpace.LOCAL, element_layout=STMatrixLayout[BM, BN, num_threads=num_threads, accum_type_size=TMemTile[dtype_, BM, BN].dtype_size].element_layout])
store_async[src_type: DType](self, src: LayoutTensor[src_type, Layout.row_major(BN), MutAnyOrigin, address_space=AddressSpace.LOCAL])
store
store[*, num_threads: Int](self, src: LayoutTensor[TMemTile[dtype_, BM, BN].dtype, STMatrixLayout[BM, BN, num_threads=num_threads, accum_type_size=TMemTile[dtype_, BM, BN].dtype_size].vec_local_layout, MutAnyOrigin, address_space=AddressSpace.LOCAL, element_layout=STMatrixLayout[BM, BN, num_threads=num_threads, accum_type_size=TMemTile[dtype_, BM, BN].dtype_size].element_layout])
store[src_type: DType](self, src: LayoutTensor[src_type, Layout.row_major(BN), MutAnyOrigin, address_space=AddressSpace.LOCAL])
load_async_with_st_matrix_layout
load_async_with_st_matrix_layout[*, num_threads: Int](self) -> LayoutTensor[TMemTile[dtype_, BM, BN].dtype, STMatrixLayout[BM, BN, num_threads=num_threads, accum_type_size=TMemTile[dtype_, BM, BN].dtype_size].vec_local_layout, MutAnyOrigin, address_space=AddressSpace.LOCAL, element_layout=STMatrixLayout[BM, BN, num_threads=num_threads, accum_type_size=TMemTile[dtype_, BM, BN].dtype_size].element_layout]
Returns:
load_async
load_async(self) -> LayoutTensor[TMemTile[dtype_, BM, BN].dtype, Layout.row_major(BN), MutAnyOrigin, address_space=AddressSpace.LOCAL]
Returns:
Was this page helpful?
Thank you! We'll create more content like this.
Thank you for helping us improve!