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 function

output_reg_to_smem_st_matrix

def output_reg_to_smem_st_matrix[output_type: DType, accum_type: DType, num_m_mmas: Int, padded_depth: Int, o_frag_size: Int, //, BM: Int, swizzle: Swizzle, num_consumer: Int](warp_group_thread_idx: UInt32, local_warp_group_idx: UInt32, output_reg_tile: TileTensor[accum_type, Layout[*?, *?], MutAnyOrigin, address_space=AddressSpace.LOCAL], accum_smem_tile: TileTensor[output_type, Layout[*?, *?], MutAnyOrigin, address_space=AddressSpace.SHARED])

Stores output register fragments to shared memory using the stmatrix PTX instruction.

Parameters:

output_type (DType): Element type of the values stored to shared memory (inferred). Must be bf16 or f16.
accum_type (DType): Element type of the accumulator fragments held in registers (inferred).
num_m_mmas (Int): Number of MMA operations along the M dimension (inferred).
padded_depth (Int): Head dimension padded to the tensor-core alignment boundary, in elements (inferred).
o_frag_size (Int): Number of elements per output fragment produced by each MMA (inferred).
BM (Int): Tile block size in the row dimension, in elements.
swizzle (Swizzle): Swizzle layout mapping stmatrix coordinates to shared memory offsets.
num_consumer (Int): Number of consumer warp groups participating in the store.

Args:

warp_group_thread_idx (UInt32): Thread index within the warp group.
local_warp_group_idx (UInt32): Local index of this warp group among the consumer warp groups.
output_reg_tile (TileTensor[accum_type, Layout[*?, *?], MutAnyOrigin, address_space=AddressSpace.LOCAL]): Register tile holding the output fragments to store.
accum_smem_tile (TileTensor[output_type, Layout[*?, *?], MutAnyOrigin, address_space=AddressSpace.SHARED]): Shared-memory tile receiving the stored output.