IMPORTANT: To view this page as Markdown, append `.md` to the URL (e.g. /max/get-started.md). For the complete documentation index, see llms.txt.
Skip to main content
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

TMemAccumulator

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

Fields​

  • ​tmem_addr (UInt32):

Implemented traits​

AccumulatorTile, AnyType, Copyable, ImplicitlyCopyable, ImplicitlyDeletable, 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])

Models

Cookbook

Puzzles

Blog

Community

Contact