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

load_matrix_a_amd

def load_matrix_a_amd[m: Int, n: Int, k: Int](a_ptr: Pointer[Float32, _safe=True], tile_row: Int, tile_col: Int, ldm: Int) -> Float32

Loads a tile of matrix A from memory to registers for AMD FP32 tensor core operations.

Constraints:

The tile dimensions must be m=16, n=16, k=4.

Parameters:

m (Int): Number of rows in the output matrix tile.
n (Int): Number of columns in the output matrix tile.
k (Int): Inner dimension for matrix multiplication.

Args:

a_ptr (Pointer[Float32, _safe=True]): Pointer to matrix A data in memory.
tile_row (Int): Starting row index of the tile.
tile_col (Int): Starting column index of the tile.
ldm (Int): Leading dimension of matrix A (stride between rows).

Returns:

Float32: SIMD vector containing 1 FP32 value loaded from matrix A.

def load_matrix_a_amd[dtype: DType, //, m: Int, n: Int, k: Int, n_blocks: Int = Int(1)](a_ptr: Pointer[Scalar[dtype], _safe=True], tile_row: Int, tile_col: Int, ldm: Int) -> SIMD[dtype, SIMDLength(4)] where dtype.is_half_float()

Loads a tile of matrix A from memory to registers for AMD half-precision tensor core operations.

Constraints:

The tile dimensions must be m=16, n=16, k=16 and n_blocks=1 or m=4, n=4, k=4 and n_blocks=16.

Parameters:

dtype (DType): Data type of the matrix elements (float16 or bfloat16).
m (Int): Number of rows in the output matrix tile.
n (Int): Number of columns in the output matrix tile.
k (Int): Inner dimension for matrix multiplication.
n_blocks (Int): Number of blocks.

Args:

a_ptr (Pointer[Scalar[dtype], _safe=True]): Pointer to matrix A data in memory.
tile_row (Int): Starting row index of the tile.
tile_col (Int): Starting column index of the tile.
ldm (Int): Leading dimension of matrix A (stride between rows).

Returns:

SIMD[dtype, SIMDLength(4)]: SIMD vector containing 4 values loaded from matrix A.