Graph Neural Networks (GNNs) have achieved significant improvements in
various domains. Sparse Matrix-Matrix multiplication (SpMM) is a fundamental
operator in GNNs, which performs a multiplication 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). 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 sparse matrix in global memory. 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. GE-SpMM performs SpMM-like
operation on sparse matrices represented in the most common Compressed Sparse
Row (CSR) format, so it can be 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 coalesced access to
sparse matrix data. We also present the Coarse-grained Warp Merging to reduce
redundant data loading among GPU warps. Experiments on a real-world graph
dataset show that GE-SpMM achieves up to 1.41X speedup over Nvidia cuSPARSE and
up to 1.81X over GraphBLAST. We also embed GE-SpMM in GNN frameworks and get up
to 3.67X speedup over popular GNN models like GCN and GraphSAGE.

本文介绍了 GE-SpMM 方法，该方法可以在 GPU 上进行支持高通用性的稀疏矩阵加速操作，并且在真实的图像数据集上实验表明，GE-SpMM 可以在 Nvidia cuSPARSE 和 GraphBLAST 上实现高达 1.41 倍和 1.81 倍的加速，且在 GNN 模型上可以获得高达 3.67 倍的加速效果。