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
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,
Deinitable,
ImplicitlyCopyable,
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
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])