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

def load_matrix_a[m: Int, n: Int, k: Int](a_ptr: Pointer[Float32, _safe=True], tile_row: Int, tile_col: Int, ldm: Int) -> SIMD[DType.float32, SIMDLength(4)]

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

Constraints:

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

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:

SIMD[DType.float32, SIMDLength(4)]: SIMD vector containing 4 TF32 values loaded from matrix A in the required order.

def load_matrix_a[m: Int, n: Int, k: Int](a_ptr: Pointer[Float16, _safe=True], tile_row: Int, tile_col: Int, ldm: Int) -> SIMD[DType.float16, SIMDLength(4)]

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

Constraints:

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

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[Float16, _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.float16, SIMDLength(4)]: SIMD vector containing 4 FP16 values loaded from matrix A in the required order.

def load_matrix_a[m: Int, n: Int, k: Int](a_ptr: Pointer[BFloat16, _safe=True], tile_row: Int, tile_col: Int, ldm: Int) -> SIMD[DType.bfloat16, (k // Int(2))]

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

Constraints:

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

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[BFloat16, _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.bfloat16, (k // Int(2))]: SIMD vector containing k//2 BF16 values loaded from matrix A in the required order.