[论文解读] GE-SpMM: General-purpose Sparse Matrix-Matrix Multiplication on GPUs for Graph Neural Networks

Guyue Huang, Guohao Dai+2 more|arXiv (Cornell University)|Jul 7, 2020

Advanced Graph Neural Networks参考文献 27被引用 6

一句话总结

GE-SpMM 为 GPU 设计了一种通用稀疏矩阵-矩阵乘法（SpMM）内核，专为图神经网络（GNNs）而优化，采用压缩稀疏行（CSR）格式以消除预处理开销。它引入了聚合行缓存（Coalesced Row Caching）和粗粒度线程束合并（Coarse-grained Warp Merging）技术，以优化内存访问并减少冗余加载，相较于 cuSPARSE 最高提升 1.41 倍，而在 DGL 和 PyG 等 GNN 框架中实现高达 3.67 倍的端到端性能提升。

ABSTRACT

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.

研究动机与目标

• 解决现有 GPU 库（如 cuSPARSE）缺乏对通用 SpMM 类操作（例如池化）支持的问题，此类限制影响了 GNN 框架的性能。
• 通过原生支持 CSR 格式，消除 GNN 框架中昂贵的数据格式转换开销，实现无需预处理的无缝集成。
• 通过利用列方向并行性和减少线程束间的冗余数据加载，克服 GPU SpMM 中低效的内存访问模式。
• 实现对多样化 GNN 模型的高性能加速，包括使用非标准归约操作（如最大池化）的模型，这些操作在厂商库中未被原生支持。

提出的方法

• 设计一种原生支持 CSR 格式的 SpMM 内核，可直接处理 CSR 格式的稀疏矩阵，避免在 GNN 管道中进行昂贵的预处理和格式转换。
• 引入聚合行缓存（Coalesced Row Caching）技术，实现列方向的并行处理，并确保内存访问模式的聚合，从而提升全局内存带宽利用率。
• 实现粗粒度线程束合并（Coarse-grained Warp Merging）技术，通过在 GPU 线程束间重用稀疏行数据，减少冗余数据加载，最小化冗余内存读取。
• 优化线程与内存访问模式，以同时利用行方向的数据重用和列方向的并行性，适配 GPU 架构特性以应对 SpMM 工作负载。
• 通过抽象归约操作，支持通用 SpMM 类操作（如最大池化），实现灵活集成到多种 GNN 模型中。
• 将内核嵌入主流 GNN 框架（DGL 和 PyG）中，评估其在真实端到端训练工作负载中的性能增益。

实验结果

研究问题

• RQ1 能否设计一种通用 SpMM 内核，在无需预处理或格式转换的前提下，同时支持标准 SpMM 和 SpMM 类操作（如最大池化）？
• RQ2 如何在 SpMM 中优化 GPU 内存访问模式，以实现聚合访问并减少线程束间的冗余数据加载？
• RQ3 在 GNN 工作负载中，原生 CSR 格式的 SpMM 内核相较于现有高性能库（如 cuSPARSE 和 GraphBLAST）能实现多大程度的性能超越？
• RQ4 将 SpMM 内核集成到 DGL 和 PyG 等真实 GNN 框架中，可在不同模型架构和数据集上实现多大的性能提升？

主要发现

• 在真实图数据集上，GE-SpMM 相较于 Nvidia cuSPARSE 最高提升 1.41 倍，相较于 GraphBLAST 最高提升 1.81 倍。
• 仅聚合行缓存技术本身即可实现最高 1.25 倍的性能提升，通过在列方向处理期间实现高效聚合内存访问。
• 粗粒度线程束合并技术贡献了最高 1.51 倍的性能提升，通过减少 GPU 线程束间的冗余数据加载。
• 在端到端 GNN 训练中，GE-SpMM 集成到 DGL 后，CUDA 时间最高减少 3.67 倍；在 PyG 中最高减少 2.10 倍，具体取决于模型配置。
• 对于 GraphSAGE-pool 中的最大池化等 SpMM 类操作，GE-SpMM 在 Pubmed 图上相较 DGL 原生实现最高提升 6.15 倍。
• 该内核在不同 GPU（GTX 1080Ti 和 RTX 2080）上均表现高效，且在不同特征向量长度和模型深度下均保持一致的性能增益。