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

TMemOperand

struct TMemOperand[dtype: DType, num_m_mmas: Int, num_n_mmas: Int, MMA_M: Int, MMA_N: Int, MMA_K: Int, num_softmax_threads: Int]

Implements a WriteableMMAOperandDescriptor backed by tensor memory (TMEM), used as the A operand of a tensor-shared (TS) UMMA.

Parameters

dtype (DType): The element type of the operand 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.
MMA_M (Int): The M dimension of a single UMMA instruction tile.
MMA_N (Int): The N dimension of a single UMMA instruction tile.
MMA_K (Int): The K dimension of a single UMMA instruction tile.
num_softmax_threads (Int): Number of threads in the consumer warp group.

Fields

tmem_addr (UInt32):

Implemented traits

AnyType, Copyable, ImplicitlyCopyable, ImplicitlyDeletable, Movable, RegisterPassable, TrivialRegisterPassable, WriteableMMAOperandDescriptor

`comptime` members

`frag_size`

comptime frag_size = TMemOperand[dtype, num_m_mmas, num_n_mmas, MMA_M, MMA_N, MMA_K, num_softmax_threads].reg_layout.frag_size

`reg_layout`

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

`reg_tile_t`

comptime reg_tile_t = LayoutTensor[dtype, TMemOperand[dtype, num_m_mmas, num_n_mmas, MMA_M, MMA_N, MMA_K, num_softmax_threads].vec_output_layout, MutAnyOrigin, address_space=AddressSpace.LOCAL, element_layout=TMemOperand[dtype, num_m_mmas, num_n_mmas, MMA_M, MMA_N, MMA_K, num_softmax_threads].reg_layout.element_layout]

`vec_output_layout`

comptime vec_output_layout = TMemOperand[dtype, num_m_mmas, num_n_mmas, MMA_M, MMA_N, MMA_K, num_softmax_threads].reg_layout.vec_output_layout

Methods

`init`

def __init__(tmem_addr: UInt32) -> Self

`offset`

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

Returns:

UInt32

`copy_from`

def copy_from[src_type: DType, src_layout: Layout, src_element_layout: Layout, //](self, src: LayoutTensor[src_type, src_layout, MutAnyOrigin, address_space=AddressSpace.LOCAL, element_layout=src_element_layout])

`copy_to`

def copy_to[dst_type: DType, dst_layout: Layout, dst_element_layout: Layout, //](self, dst: LayoutTensor[dst_type, dst_layout, MutAnyOrigin, address_space=AddressSpace.LOCAL, element_layout=dst_element_layout])

Parameters​

Fields​

Implemented traits​

comptime members​

frag_size​

reg_layout​

reg_tile_t​

vec_output_layout​

Methods​

__init__​

offset​

copy_from​

copy_to​