Skip to main content

Mojo struct

TMemOperand

@register_passable(trivial) 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]

Fields

  • tmem_addr (UInt32):

Implemented traits

AnyType, Copyable, ImplicitlyCopyable, Movable, UnknownDestructibility, WriteableMMAOperandDescriptor

Aliases

__copyinit__is_trivial

alias __copyinit__is_trivial = UInt32.__copyinit__is_trivial

__del__is_trivial

alias __del__is_trivial = UInt32.__del__is_trivial

__moveinit__is_trivial

alias __moveinit__is_trivial = UInt32.__moveinit__is_trivial

frag_size

alias frag_size = ((MMA_M * MMA_N) // num_softmax_threads)

reg_layout

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

reg_tile_t

alias reg_tile_t = LayoutTensor[dtype, Layout.__init__(IntTuple.__init__[__origin_of()](IntTuple[__origin_of()](2, num_m_mmas), IntTuple[__origin_of()]((((MMA_M * MMA_N) // num_softmax_threads) // 4), num_n_mmas), Tuple[]()), IntTuple.__init__[__origin_of()](IntTuple[__origin_of()](2, ((MMA_M * MMA_N) // num_softmax_threads)), IntTuple[__origin_of()](4, (((MMA_M * MMA_N) // num_softmax_threads) * num_m_mmas)), Tuple[]())), MutableAnyOrigin, address_space=AddressSpace(5), element_layout=Layout.row_major(1, 2)]

vec_output_layout

alias vec_output_layout = Layout.__init__(IntTuple.__init__[__origin_of()](IntTuple[__origin_of()](2, num_m_mmas), IntTuple[__origin_of()]((((MMA_M * MMA_N) // num_softmax_threads) // 4), num_n_mmas), Tuple[]()), IntTuple.__init__[__origin_of()](IntTuple[__origin_of()](2, ((MMA_M * MMA_N) // num_softmax_threads)), IntTuple[__origin_of()](4, (((MMA_M * MMA_N) // num_softmax_threads) * num_m_mmas)), Tuple[]()))

Methods

__init__

__init__(tmem_addr: UInt32) -> Self

offset

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

Returns:

UInt32

copy_from

copy_from[src_type: DType, src_layout: Layout, src_element_layout: Layout, //](self, src: LayoutTensor[src_type, src_layout, MutableAnyOrigin, address_space=AddressSpace(5), element_layout=src_element_layout])

copy_to

copy_to[dst_type: DType, dst_layout: Layout, dst_element_layout: Layout, //](self, dst: LayoutTensor[dst_type, dst_layout, MutableAnyOrigin, address_space=AddressSpace(5), element_layout=dst_element_layout])

Was this page helpful?