[论文解读] Accelerating SpMM Kernel with Cache-First Edge Sampling for Graph Neural Networks

Chien‐Yu Lin, Liang Luo+1 more|arXiv (Cornell University)|Apr 21, 2021

Advanced Graph Neural Networks参考文献 30被引用 5

一句话总结

该论文提出 ES-SpMM，一种以缓存为中心的边缘采样机制及其协同设计的 SpMM 内核，通过采样边将图数据适配至 GPU 共享内存，从而减少计算量并提升缓存局部性，加速 GNN 推理。在标准 GNN 基准测试中，该方法在不损失准确率的情况下实现最高 4.35 倍于 cuSPARSE 的加速，而在准确率损失小于 1% 的情况下实现高达 45.3 倍的加速。

ABSTRACT

Graph neural networks (GNNs), an emerging deep learning model class, can extract meaningful representations from highly expressive graph-structured data and are therefore gaining popularity for wider ranges of applications. However, current GNNs suffer from the poor performance of their sparse-dense matrix multiplication (SpMM) operator, even when using powerful GPUs. Our analysis shows that 95% of the inference time could be spent on SpMM when running popular GNN models on NVIDIA's advanced V100 GPU. Such SpMM performance bottleneck hinders GNNs' applicability to large-scale problems or the development of more sophisticated GNN models. To address this inference time bottleneck, we introduce ES-SpMM, a cache-first edge sampling mechanism and codesigned SpMM kernel. ES-SpMM uses edge sampling to downsize the graph to fit into GPU's shared memory. It thus reduces the computation cost and improves SpMM's cache locality. To evaluate ES-SpMM's performance, we integrated it with a popular GNN framework, DGL, and tested it using representative GNN models and datasets. Our results show that ES-SpMM outperforms the highly optimized cuSPARSE SpMM kernel by up to 4.35x with no accuracy loss and by 45.3x with less than a 1% accuracy loss.

研究动机与目标

• 解决 GPU 上稀疏-密集矩阵乘法（SpMM）导致的 GNN 推理性能瓶颈问题。
• 在不牺牲模型准确率的前提下，减少 SpMM 计算时间并提升缓存局部性。
• 通过将边缘采样直接集成到 SpMM 内核中，消除预处理开销。
• 实现与现有 GNN 框架（如 DGL）的无缝集成，无需修改用户代码。
• 提供一种通用、高性能的 SpMM 解决方案，适用于多种 GNN 工作负载。

提出的方法

• 提出一种内核内边缘采样机制，在 GPU 上执行采样，消除 CPU 预处理步骤，并利用 GPU 的并行性。
• 采用以缓存为中心的设计，将共享内存大小作为采样上限，确保采样后的图完全容纳于共享内存中。
• 使用合并的全局内存访问模式，高效地将采样数据加载到共享内存中。
• 采用两种轻量级采样策略——Bucket 和 FastRand，根据数据集特性在速度与准确率之间实现平衡。
• 支持标准 CSR 格式，避免格式转换开销，并可在不修改用户代码的情况下作为 DGL 的即插即用替代方案。
• 将采样逻辑与 SpMM 内核协同设计，以统一优化内存访问与计算量减少。

实验结果

研究问题

• RQ1 是否可将边缘采样集成到 SpMM 内核中，以消除预处理开销并加速推理？
• RQ2 是否可通过以缓存为中心的采样策略（将采样图适配至共享内存）显著提升 SpMM 性能？
• RQ3 在使用内核内边缘采样时，SpMM 性能可实现多大程度的加速，且准确率损失最小？
• RQ4 与高度优化的 SpMM 内核（如 cuSPARSE）相比，ES-SpMM 在速度与准确率方面表现如何？
• RQ5 ES-SpMM 是否可无缝集成到现有 GNN 框架中，而无需修改用户级代码？

主要发现

• 在标准 GNN 基准测试中，ES-SpMM 在不损失准确率的情况下，实现最高 4.35 倍于高度优化的 cuSPARSE SpMM 内核的加速。
• 在准确率损失小于 1% 的情况下，ES-SpMM 实现高达 45.3 倍于 cuSPARSE 的加速，展现出卓越的性能-准确率权衡。
• 在标准 GCN 模型中，高达 95% 的推理时间消耗在 SpMM 上，证实 SpMM 是主要性能瓶颈。
• 内核内采样机制相比 CPU 预处理减少了开销，实现了更快、更高效的采样。
• ES-SpMM 在多种 GNN 模型（GCN 和 GraphSAGE）和数据集（Reddit、Cora、PubMed、Ppi）上均保持高性能。
• 与 DGL 的集成透明，无需修改用户级代码，且支持标准 CSR 格式，无需格式转换。