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
TMemTile
struct TMemTile[dtype_: DType, BM: Int, BN: Int]
Represents a tile in SM100 tensor memory (TMEM) and provides async load/store helpers.
Parameters
- dtype_ (
DType): Element dtype of the TMEM tile. - BM (
Int): Number of rows in the tile, in elements; must be a multiple of 64. - BN (
Int): Number of columns in the tile, in elements.
Fields
- tmem_addr (
UInt32):
Implemented traits
AnyType,
Copyable,
ImplicitlyCopyable,
ImplicitlyDeletable,
Movable,
RegisterPassable,
TrivialRegisterPassable
comptime members
dtype
comptime dtype = dtype_
dtype_size
comptime dtype_size = size_of[TMemTile[dtype_, BM, BN].dtype]()
num_m_tiles
comptime num_m_tiles = (BM // Int(64))
Methods
__init__
def __init__(tmem_addr: UInt32) -> Self
__getitem__
def __getitem__(self, i: UInt32) -> Self
offset
allocate_register_tile
static def allocate_register_tile[*, num_threads: Int]() -> LayoutTensor[Self.dtype, STMatrixLayout[BM, BN, num_threads=num_threads, accum_dtype_size=size_of[dtype_]()].vec_local_layout, MutAnyOrigin, address_space=AddressSpace.LOCAL, element_layout=STMatrixLayout[BM, BN, num_threads=num_threads, accum_dtype_size=size_of[dtype_]()].element_layout]
Returns:
store_async
def store_async[*, num_threads: Int](self, src: LayoutTensor[Self.dtype, STMatrixLayout[BM, BN, num_threads=num_threads, accum_dtype_size=size_of[dtype_]()].vec_local_layout, MutAnyOrigin, address_space=AddressSpace.LOCAL, element_layout=STMatrixLayout[BM, BN, num_threads=num_threads, accum_dtype_size=size_of[dtype_]()].element_layout])
def store_async[src_type: DType](self, src: TileTensor[src_type, Layout[*?, *?], MutUntrackedOrigin, address_space=AddressSpace.LOCAL])
def store_async[src_type: DType, src_len: Int, src_offset: Int = Int(0)](self, src: Array[Scalar[src_type], src_len])
load_async_with_st_matrix_layout
def load_async_with_st_matrix_layout[*, num_threads: Int](self) -> LayoutTensor[Self.dtype, STMatrixLayout[BM, BN, num_threads=num_threads, accum_dtype_size=size_of[dtype_]()].vec_local_layout, MutAnyOrigin, address_space=AddressSpace.LOCAL, element_layout=STMatrixLayout[BM, BN, num_threads=num_threads, accum_dtype_size=size_of[dtype_]()].element_layout]
Returns:
load_st_matrix_chunk
def load_st_matrix_chunk[*, num_threads: Int, start_repeat: Int, num_repeats: Int](self, dst: LayoutTensor[Self.dtype, STMatrixLayout[BM, BN, num_threads=num_threads, accum_dtype_size=size_of[dtype_]()].vec_local_layout, MutAnyOrigin, address_space=AddressSpace.LOCAL, element_layout=STMatrixLayout[BM, BN, num_threads=num_threads, accum_dtype_size=size_of[dtype_]()].element_layout])
Load a range of repeat columns from tmem into a pre-allocated tensor.
Parameters:
- num_threads (
Int): Number of threads in the warp group. - start_repeat (
Int): First repeat index to load (0-based). - num_repeats (
Int): Number of repeats to load.
Args:
- dst (
LayoutTensor[Self.dtype, STMatrixLayout[BM, BN, num_threads=num_threads, accum_dtype_size=size_of[dtype_]()].vec_local_layout, MutAnyOrigin, address_space=AddressSpace.LOCAL, element_layout=STMatrixLayout[BM, BN, num_threads=num_threads, accum_dtype_size=size_of[dtype_]()].element_layout]): Pre-allocated register tensor.
load_async
Was this page helpful?
Thank you! We'll create more content like this.
Thank you for helping us improve!