GE-SpMM: general-purpose sparse matrix-matrix multiplication on GPUs for graph neural networks

Graph Neural Networks (GNNs) based graph learning algorithms have achieved significant improvements in various domains. Sparse Matrix-Matrix multiplication (SpMM) is a fundamental operator in GNNs, which performs a multiplication operation between a sparse matrix and a dense matrix. Accelerating SpMM on parallel hardware like GPUs can face the following challenges: From the GNN application perspective, the compatibility needs to be considered. General GNN algorithms require SpMM-like operations (e.g., pooling) between matrices, which are not supported in current high-performance GPU libraries (e.g., Nvidia cuSPARSE [1]). Moreover, the sophisticated preprocessing in previous implementations will lead to heavy data format conversion overheads in GNN frameworks. From the GPU hardware perspective, optimizations in SpMV (Sparse Matrix-Vector) designs on GPUs do not apply well to SpMM. SpMM exposes the column-wise parallelism in the dense output matrix, but straightforward generalization from SpMV leads to inefficient, uncoalesced access to the GPU global memory. Moreover, the sparse row data can be reused among GPU threads, which is neither possible in SpMM designs inherited from SpMV.

To tackle these challenges, we propose GE-SpMM 1 . GE-SpMM performs SpMM-like operation on sparse matrices represented in the most common Compressed Sparse Row (CSR) format. Thus, GE-SpMM can be efficiently embedded in GNN frameworks with no preprocessing overheads and support general GNN algorithms. We introduce the Coalesced Row Caching method to process columns in parallel and ensure efficient coalesced access to the GPU global memory. We also present the Coarsegrained Warp Merging method to reduce redundant data loading among GPU warps. Experiments on a real-world graph dataset show that GE-SpMM achieves up to 1.41× speedup over Nvidia cuSPARSE [1] and up to 1.81× over GraphBLAST [2]. We also embed GE-SpMM in GNN frameworks and get up to 3.67× speedup over popular GNN models like GCN [3] and GraphSAGE [4].