std :: ml::gpu_mps
Function std::mps_matmul_f32
int mps_matmul_f32(
[]const float a,
[]const float b,
[]float out,
unsigned long M,
unsigned long K,
unsigned long N,
)out[M,N] = a[M,K] . b[K,N], computed on the GPU. float32-only (Metal has no 'double' type at all). Internally picks between three kernels depending on shape, most-specialized first:
- M, K, N all exact multiples of 32: matmul_kernel_smma_multi --
threadgroup-memory tiling (data reuse, like matmul_kernel) AND
simdgroup_matrix hardware matrix units (like matmul_kernel_smma)
together, the same combination PyTorch's MPSMatrixMultiplication and
MLX's fastest GEMM paths use.
- otherwise, M, K, N all exact multiples of 8: matmul_kernel_smma --
simdgroup_matrix alone, no threadgroup-memory tiling.
- otherwise: matmul_kernel's ordinary threadgroup-memory tiling, no
alignment requirement at all.
The two simdgroup_matrix kernels require their alignment because simdgroup_load/store past a matrix's real bounds is undefined behavior, unlike matmul_kernel's explicit boundary checks. See gpu_mps.sc's comment on the implementation, and gpu_mps_kernels.metal's comments on each kernel, for detail on all three paths.