For the complete documentation index, see llms.txt. Markdown versions of all pages are available by appending .md to any URL (e.g. /get-started.md).

Mojo struct

TMemAccumulator

struct TMemAccumulator[dtype_: DType, MMA_M: Int, MMA_N: Int, num_m_mmas: Int, num_n_mmas: Int, num_softmax_threads: Int]

Implements an AccumulatorTile backed by tensor memory (TMEM), storing UMMA results at a given TMEM address.

Parameters

• ​dtype_ (DType): The element type of the accumulator tile.
• ​MMA_M (Int): The M dimension of a single UMMA instruction tile.
• ​MMA_N (Int): The N dimension of a single UMMA instruction tile.
• ​num_m_mmas (Int): Number of UMMA tiles along the M dimension.
• ​num_n_mmas (Int): Number of UMMA tiles along the N dimension.
• ​num_softmax_threads (Int): Number of threads in the consumer warp group.

Fields

• ​tmem_addr (UInt32):

Implemented traits

AccumulatorTile, AnyType, Copyable, Deinitable, ImplicitlyCopyable, Movable, RegisterPassable, TrivialRegisterPassable

comptime members

dtype

comptime dtype = dtype_

element_layout

comptime element_layout = TMemAccumulator[dtype_, MMA_M, MMA_N, num_m_mmas, num_n_mmas, num_softmax_threads].layout_t.element_layout

frag_size

comptime frag_size = TMemAccumulator[dtype_, MMA_M, MMA_N, num_m_mmas, num_n_mmas, num_softmax_threads].layout_t.frag_size

layout_t

comptime layout_t = RegisterAccumulatorLayout[MMA_M, MMA_N, num_m_mmas, num_n_mmas, num_softmax_threads]

rows_of_frags_layout

comptime rows_of_frags_layout = TMemAccumulator[dtype_, MMA_M, MMA_N, num_m_mmas, num_n_mmas, num_softmax_threads].layout_t.rows_of_frags_layout

vec_output_layout

comptime vec_output_layout = TMemAccumulator[dtype_, MMA_M, MMA_N, num_m_mmas, num_n_mmas, num_softmax_threads].layout_t.vec_output_layout

Methods

__init__

def __init__(tmem_addr: UInt32) -> Self

__getitem__

def __getitem__(self, i: UInt32) -> Self

check_constraints

static def check_constraints()

offset

def offset[m_mma: Int, n_mma: Int](self) -> UInt32

Returns:

UInt32

rows_of_frags

static def rows_of_frags(src: LayoutTensor[Self.dtype, Self.vec_output_layout, MutAnyOrigin, address_space=AddressSpace.LOCAL, element_layout=Self.layout_t.element_layout]) -> LayoutTensor[Self.dtype, Self.rows_of_frags_layout, MutAnyOrigin, address_space=AddressSpace.LOCAL]

Returns:

LayoutTensor[Self.dtype, Self.rows_of_frags_layout, MutAnyOrigin, address_space=AddressSpace.LOCAL]

allocate_register_tile

static def allocate_register_tile() -> LayoutTensor[Self.dtype, Self.vec_output_layout, MutAnyOrigin, address_space=AddressSpace.LOCAL, element_layout=Self.layout_t.element_layout]

Returns:

LayoutTensor[Self.dtype, Self.vec_output_layout, MutAnyOrigin, address_space=AddressSpace.LOCAL, element_layout=Self.layout_t.element_layout]

copy_from

def copy_from(self, src: LayoutTensor[Self.dtype, Self.vec_output_layout, MutAnyOrigin, address_space=AddressSpace.LOCAL, element_layout=Self.layout_t.element_layout])

copy_to

def copy_to(self, dst: LayoutTensor[Self.dtype, Self.vec_output_layout, MutAnyOrigin, address_space=AddressSpace.LOCAL, element_layout=Self.layout_t.element_layout])