TileGEMM: Boosting the Performance of GEMM on AMX-Powered CPUs by Exploiting Data Reuse

Abstract

General Matrix Multiplication (GEMM) is the cornerstone of high-performance computing and deep learning. Its efficiency significantly influences the performance of applications ranging from large language models to scientific simulations. Intel Advanced Matrix Extensions (AMX) significantly boost matrix operations throughput, yet existing implementations underutilize its potential. The core bottleneck is that AMX’s unique microarchitectural constraints break the classic Advanced Vector Extensions 512 (AVX-512) GEMM paradigm, causing critical data reuse loss across the memory hierarchy: inefficient vector register reuse during mandatory layout transformation, poorly reconciled trade-off between Tile register reuse for partial sums and L1 cache locality, and severe buffer reuse degradation due to fine-grained update strategies. To address these challenges, we propose TileGEMM, a high-performance GEMM implementation on AMX that systematically enhances data reuse across the memory hierarchy. First, a dual Tile register layout transformation strategy eliminates redundant memory accesses, maximizing vector register reuse. Second, a multi-level cache- and Tile register-aware tiling strategy adopts a K-dimension-centric model under strict L1 capacity constraints, improving Tile register reuse for intermediate partial sums while preserving L1 cache locality. Third, a coarse-grained buffer update strategy with loop reordering extends transformed data residency, boosting buffer reuse and reducing transformation overhead. Extensive experiments in various matrix configurations show that TileGEMM achieves average speedups of 3.27 × and 1.27 × over AMX-based implementations LIBXSMM and OneDNN, and 10.30 × and 7.96 × over AVX-512-based implementations TVM and MKL, respectively.