A Machine Learning Approach Towards Runtime Optimisation of Matrix Multiplication

TL;DR

The paper tackles runtime optimization of GEMM on multi-core shared-memory systems by automatically selecting the optimal thread count using architecture-aware machine learning. It introduces ADSALA, a library that collects GEMM timing data at installation, trains an ML predictor (notably XGBoost) on features derived from matrix sizes and parallelism, and uses the predictor at runtime to minimize GEMM time with an overhead that remains small. Results on two HPC platforms ( Cascade Lake and Zen 3) show average speedups of 25–40% for matmuls with memory footprints under 100 MB, demonstrating that learned, platform-specific threading policies can outperform conventional fixed-core strategies. The work highlights the practical viability of architecture-aware BLAS optimization and lays groundwork for extending to additional BLAS routines and heterogeneous architectures.

Abstract

The GEneral Matrix Multiplication (GEMM) is one of the essential algorithms in scientific computing. Single-thread GEMM implementations are well-optimised with techniques like blocking and autotuning. However, due to the complexity of modern multi-core shared memory systems, it is challenging to determine the number of threads that minimises the multi-thread GEMM runtime. We present a proof-of-concept approach to building an Architecture and Data-Structure Aware Linear Algebra (ADSALA) software library that uses machine learning to optimise the runtime performance of BLAS routines. More specifically, our method uses a machine learning model on-the-fly to automatically select the optimal number of threads for a given GEMM task based on the collected training data. Test results on two different HPC node architectures, one based on a two-socket Intel Cascade Lake and the other on a two-socket AMD Zen 3, revealed a 25 to 40 per cent speedup compared to traditional GEMM implementations in BLAS when using GEMM of memory usage within 100 MB.

Figures (13)

• Figure 1: Histogram of the optimal thread count; measured with single precision GEMM with memory usage within 100 MB. The experiment was run on a single HPC node with 2 Cascade Lake CPUs and 48 physical cores. the BLAS package used is Math Kernel Library.
• Figure 2: The installation workflow of ADSALA GEMM. Upon ADSALA installation, the library performs two sub-parts shown in the diagram. In the end, two files containing the configurations together with the production-ready ML model will be saved for later use at runtime.
• Figure 3: The runtime workflow of ADSALA GEMM. Configuration file and trained ML model output during installation (see Fig. \ref{['fig:Lib Structure']}) are used by the runtime library.
• Figure 4: Feature distributions before and after transformation. The data is from GEMM with input matrix size below 500 MB on Setonix.
• Figure 5: A schematic diagram for the 2-socket EPYC CPU configuration on Setonix.