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

matmul_kernel

def matmul_kernel[c_type: DType, a_type: DType, b_type: DType, tile_size: Int, elementwise_lambda_fn: Optional[def[dtype: DType, width: SIMDLength, *, alignment: Int = Int(1)](IndexList[Int(2)], SIMD[dtype, width]) capturing thin -> None] = None, s_type: DType = get_accum_type[c_type]()](c_ptr: Pointer[Scalar[c_type], MutAnyOrigin, _safe=False], a_ptr: Pointer[Scalar[a_type], ImmutAnyOrigin, _safe=False], b_ptr: Pointer[Scalar[b_type], ImmutAnyOrigin, _safe=False], m: Int, n: Int, k: Int)

Matrix Multiplication using shared memory. This version loads blocks of size tile_size x tile_size from A and B and updates a tile_size x tile_size in C. The thread block should have shape (tile_size, tile_size, 1). Each thread is mapped one element in C. The grid should have shape (N/tile_size, M/tile_size, 1). N is the first dimension for coalesced access.

Parameters:

c_type (DType): DType of the output matrix C elements.
a_type (DType): DType of the input matrix A elements.
b_type (DType): DType of the input matrix B elements.
tile_size (Int): Side length, in elements, of the square shared-memory tile loaded from A and B and written to C.
elementwise_lambda_fn (Optional[def[dtype: DType, width: SIMDLength, *, alignment: Int = Int(1)](IndexList[Int(2)], SIMD[dtype, width]) capturing thin -> None]): Optional epilogue applied to the accumulated value before it is stored to C (defaults to None, which stores the raw accumulation).
s_type (DType): DType used for the inner accumulation (defaults to the accumulator type for c_type).

Args:

c_ptr (Pointer[Scalar[c_type], MutAnyOrigin, _safe=False]): Pointer to the row-major output matrix C of shape (m, n).
a_ptr (Pointer[Scalar[a_type], ImmutAnyOrigin, _safe=False]): Pointer to the row-major input matrix A of shape (m, k).
b_ptr (Pointer[Scalar[b_type], ImmutAnyOrigin, _safe=False]): Pointer to the row-major input matrix B of shape (k, n).
m (Int): Number of rows of A and C.
n (Int): Number of columns of B and C.
k (Int): Contraction dimension: columns of A and rows of B.