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 function

store_matrix_d

def store_matrix_d[dtype: DType, //, m: Int, n: Int, k: Int, n_blocks: Int = Int(1)](d_ptr: Pointer[Scalar[dtype], _safe=True], d: SIMD[dtype, SIMDLength(4)], tile_row: Int, tile_col: Int, ldm: Int)

Stores matrix D tile from registers to memory after tensor core operation.

This function dispatches to architecture-specific implementations for storing the results of a tensor core matrix multiply-accumulate operation. It handles the different memory layouts required by NVIDIA and AMD tensor cores.

Note:

Automatically selects appropriate implementation based on GPU architecture.
Each thread stores 4 elements in architecture-specific positions.
Must be called by all threads in a warp.

Parameters:

dtype (DType): Data type of the matrix elements.
m (Int): Number of rows in matrix D.
n (Int): Number of columns in matrix D.
k (Int): Inner dimension for matrix multiply.
n_blocks (Int): Number of blocks.

Args:

d_ptr (Pointer[Scalar[dtype], _safe=True]): Pointer to destination memory for matrix D.
d (SIMD[dtype, SIMDLength(4)]): SIMD vector containing 4 elements to store.
tile_row (Int): Starting row index of the tile in matrix D.
tile_col (Int): Starting column index of the tile in matrix D.
ldm (Int): Leading dimension (stride) of matrix D.