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

stsm_helper

def stsm_helper[swizzle: Swizzle, vec_dtype: DType, vec_size: Int, transpose_c: Bool = False](vec: Array[Scalar[vec_dtype], vec_size], dst: LayoutTensor[address_space=AddressSpace.SHARED, element_layout=dst.element_layout, layout_int_type=dst.layout_int_type, linear_idx_type=dst.linear_idx_type, masked=dst.masked, alignment=dst.alignment])

Stores a register fragment to shared memory using the stmatrix instruction.

Casts the loaded fragment to the destination dtype, applies the shared memory swizzle, and emits st_matrix stores so the epilogue can drain tensor memory through shared memory before the TMA store.

Parameters:

swizzle (Swizzle): Shared-memory swizzle applied to compute the store offset from the lane offset.
vec_dtype (DType): Element type of the input register fragment vec.
vec_size (Int): Number of scalar elements in the input register fragment vec.
transpose_c (Bool): Whether to transpose the fragment layout when computing the store offset (defaults to False).

Args:

vec (Array[Scalar[vec_dtype], vec_size]): Register fragment loaded from tensor memory to store to shared memory.
dst (LayoutTensor[address_space=AddressSpace.SHARED, element_layout=dst.element_layout, layout_int_type=dst.layout_int_type, linear_idx_type=dst.linear_idx_type, masked=dst.masked, alignment=dst.alignment]): Destination shared-memory LayoutTensor where the fragment is written via st_matrix.