LIBSHALOM: Optimizing Small and Irregular-Shaped Matrix Multiplications on ARMv8 Multi-Core Processors

The International Conference for High Performance Computing, Networking, Storage and Analysis (SC) was held from November 14-19, 2021, in St. Louis, Missouri, USA. SC is one of the top conferences in the field of high-performance computing, primarily featuring significant academic contributions in areas such as system architecture, high-performance networking, and system performance evaluation. To date, it has been held 32 times, with an average acceptance rate of 22%. This year, 379 papers were submitted, of which 98 were accepted, resulting in an acceptance rate of 25.9%.

A master’s student from the High-Performance Networking and Architecture (HiNA) research group at the College of Computer Science, National University of Defense Technology, Yang Weiling, under the guidance of professors Dong Dezun and Fang Jianbin, and in collaboration with Wang Zheng and Su Xing, has authored the academic paper “LibShalom: Optimizing Small and Irregular-Shaped Matrix Multiplications on ARMv8 Multi-Cores,” which has been accepted and published at this conference. This paper is also the first publication from the HiNA research group at SC.

The paper investigates performance optimization techniques for general matrix multiplication (GEMM) of small and irregular shapes based on ARMv8 multi-core processors, leading to the design and implementation of the LibShalom high-performance computing library. GEMM is one of the most widely used high-performance computing kernels across various fields, including deep learning, signal processing, advanced physical analysis, and astrophysics. GEMM refers to the multiplication of dense matrices, expressed as: C = αA×B + βC, where A, B, and C are matrices of sizes M×K, K×N, and M×N respectively, and α and β are scalars. With the development of the HPC field, the workloads of GEMM in many practical applications, such as CP2K and CFD, exhibit small and irregular characteristics. There is still significant optimization potential for these types of GEMM on ARMv8 multi-core processors. Therefore, this paper improves upon traditional GEMM algorithms by proposing to overlap packing and computation, implemented through hand-written assembly code, and reasonably distributing the overall GEMM workload across each computational core of the processor, effectively enhancing the performance of small and irregular GEMM. This achievement is widely applicable to scientific computing and intelligent computing applications targeting ARMv8 processors.

Research Motivation and Progress

Linear algebra kernels are widely used in scientific research and engineering computations. GEMM is a key operation in linear algebra and serves as a fundamental building block for HPC applications—from traditional scientific simulations to emerging deep learning workloads. Although GEMM optimization is a well-researched area, existing linear algebra libraries (BLAS) primarily focus on optimizing large matrices of regular shapes.

LIBSHALOM: Optimizing Small and Irregular-Shaped Matrix Multiplications on ARMv8 Multi-Core Processors

Figure 1: Usage scenarios of GEMM in real programs

Due to the diversity and continuous evolution of HPC workloads, the sizes and shapes of input matrices in GEMM can vary based on the algorithms and input data of applications. For instance, in applications like CFD and CP2K, commonly used matrix sizes are 5 × 5 and 23 × 23. Additionally, some data analysis algorithms require operations on irregularly shaped matrices, where the sizes of the two dimensions of the matrix differ significantly. For example, the convolution calculations used in the ResNet neural network employ GEMM, where one dimension of the matrix is 64 while the other dimension exceeds 3000.

To identify the performance bottlenecks of existing BLAS libraries for small and irregular matrix operations on ARMv8 processors, the paper evaluates the GEMM performance of four BLAS libraries: OpenBLAS, BLIS, ARMPL, and BLASFEO on the Phytium 2000+ processor. Figure 2 (a) shows the performance of square matrix GEMM, from which we observe that existing BLAS libraries perform poorly for smaller matrix sizes. For example, when the matrix size is 32, even the best-performing BLASFEO only achieves about 60% of peak performance. Figure 2 (b) illustrates the performance of irregularly shaped GEMM; for this type of matrix, the highly optimized BLIS can achieve 70% of peak performance when M is 4096. However, when M is less than 128, all libraries fail to reach 40% of the theoretical peak performance.

The performance bottlenecks of existing BLAS for small and irregular-shaped GEMM on ARMv8 processors mainly lie in: (1) the packing overhead occupies too large a proportion of the GEMM runtime; (2) inefficient handling of the edges for GEMM; (3) unreasonable parallelization for irregularly shaped GEMM. This paper aims to address the shortcomings of existing BLAS libraries, significantly improving the performance of small and irregular GEMM on ARMv8 processors.

LIBSHALOM: Optimizing Small and Irregular-Shaped Matrix Multiplications on ARMv8 Multi-Core Processors

Figure 2: Performance of existing BLAS libraries on Phytium 2000+

Technical Details

This paper proposes to integrate the packing step of GEMM into the computation to hide the packing overhead. During the GEMM computation process, the processor executes many fused multiply-add (FMA) instructions, during which the memory access components are idle. By inserting the memory access instructions required for packing into the FMA instructions, the memory access components can remain busy, effectively overlapping computation with packing. To handle the edges of GEMM, this paper also designs an efficient edge microkernel to hide the overhead of data access. Existing BLAS libraries, such as OpenBLAS, have not fully optimized the scheduling of instructions in their edge microkernels, with dependencies between read-after-write instructions being too close together, failing to hide the latency of data access instructions. LibShalom rearranges the dependent instructions in the edge microkernel, placing the data access instructions as early as possible to ensure that data is already fetched into vector registers when used. Additionally, this paper employs a two-level parallelization strategy to parallelize irregularly shaped GEMM, dividing the M and N dimensions of the task matrix C, with the number of threads used for each dimension determined by the compute-to-memory access ratio.

The overall algorithm flow of LibShalom is shown in Figure 3, which is for the NN mode of GEMM; other modes (e.g., NT) can also be applied.

LIBSHALOM: Optimizing Small and Irregular-Shaped Matrix Multiplications on ARMv8 Multi-Core Processors

Figure 3

This research focuses on optimizing small and irregular-shaped GEMM on ARMv8 processors, requiring evaluation and validation across multiple ARMv8 processors. The experiments in the paper compare LibShalom with five GEMM libraries on three ARMv8 processors. The five libraries are: OpenBLAS, BLIS, ARMPL, BLASFEO, and LIBXSMM. The performance parameters of the three ARM processors are shown in the following table:

LIBSHALOM: Optimizing Small and Irregular-Shaped Matrix Multiplications on ARMv8 Multi-Core Processors

The paper first evaluates the performance of single-threaded small GEMM, taking the average of 10 runs, with the results shown in Figure 4. In this case, LibShalom consistently outperforms other methods. The advantage of LibShalom is particularly evident with smaller matrices; for instance, when M = N = K = 8, the throughput of LibShalom is twice that of BLASFEO. Overall, LibShalom achieves a speedup of 1.05 to 2 times for small GEMM across different ARMv8 hardware platforms.

LIBSHALOM: Optimizing Small and Irregular-Shaped Matrix Multiplications on ARMv8 Multi-Core Processors

Figure 4

For multi-threaded irregular GEMM, the paper also conducted performance evaluations. Figure 5 shows the performance of GEMM in NT mode on the Phytium 2000+, while Figure 6 displays the performance of GEMM in NN and NT modes on the KP920 and ThunderX2 platforms. Overall, LibShalom is universally applicable, achieving high performance across different ARMv8 processors, with an average speedup of 1.4 times.

LIBSHALOM: Optimizing Small and Irregular-Shaped Matrix Multiplications on ARMv8 Multi-Core Processors

Figure 5

LIBSHALOM: Optimizing Small and Irregular-Shaped Matrix Multiplications on ARMv8 Multi-Core Processors

Figure 6

Paper link:

https://dl.acm.org/doi/10.1145/3458817.3476217

Code link:

https://github.com/AnonymousYWL/LibShalom

We welcome everyone to try it out and star the repository.

Leave a Comment