Abstract:

General matrix multiplication (GEMM) is the most important operation in linear algebra, and many applications from different scientific fields have converted their key parts to use GEMM. GEMM is widely used in large models, machine learning, scientific computing, signal processing, and other fields. In particular, half-precision batch GEMM (i.e., FP16) has been the core operation of many deep learning frameworks. This paper proposes a GPU optimization scheme for half-precision batch GEMM (HGEMM). In terms of blocking strategy, this scheme provides a matrix size-affinitive blocking strategy to ensure that each wavefront is assigned the same workload and computation; threads calculate multiple matrix multiplications simultaneously to improve the utilization of computing units. In terms of memory access optimization: at the cost of multiple data reads, the same amount of memory access is allocated to each thread to facilitate compiler optimization, ensuring that memory access and computation time cover each other. For extremely small-size batch HGEMM with a matrix size of less than 16, this paper uses a 4x4x4 Matrix Core and its corresponding blocking scheme to improve memory access performance while reducing computation waste, and provides an option to use shared memory to achieve the highest performance. This paper compares the performance of this scheme with two operators of rocBLAS on the AMD GPU MI210 platform. The results show that the average performance of this scheme on AMD GPU MI210 is 4.14 times that of rocBLAS hgemm batched and 4.96 times that of rocBLAS gemm ex batched. In particular, for extremely small batch sizes, the average performance of HGEMM is 18.60 times that of rocBLAS hgemm batched and 14.02 times that of rocBLAS gemm ex batched.

Key words: GPU, matrix core, matrix multiplication, memory access optimization