Accelerating Batched Matrix Multiplication for Variable Small Sizes Based on TVM andApplications

In many practical applications,efficient computation of a large amount of small matrix products across different dimensions is required.For instance,in graph classification tasks based on graph neural networks,multiple adjacency matrices need to be multiplied with node feature matrices.To address the issue of existing methods being unable to efficiently handle batched matrix multiplication for variable and small sizes across different hardware platforms,this paper introduces a cross-platform efficient algorithm,BVSM,based on the deep learning compiler TVM.BVSM enhances the capability of TVM to efficiently perform batched matrix multiplication for small and variable sizes by utilizing customized optimization templates for small matrices,employing tensorization for batching,and applying grouped padding techniques.Experiments on real datasets of graph classification demonstrate that,on the CPU,BVSM achieves on average over two times speedup compared to auto-scheduled and auto-tuned TVM(AnsorTVM),reaching 95% of the average performance of,and achieving up to 1.27 times compared to,the method of batched matrix multiplication for variable sizes of Intel MKL.On the GPU,BVSM achieves on average over 62.05 times speedup compared to AnsorTVM,over 28.82 times speedup compared to cuBLAS,and over 6.59 times speedup compared to the method of batched matrix multiplication for variable sizes of MAGMA.