AES-SpMM：基于自适应边采样策略加速GNN中稀疏-稠密矩阵乘法
分布式、并行与集群计算
2025-03-25 v1
摘要
在加速图神经网络(GNN)方面，稀疏-稠密矩阵乘法(SpMM)的采样策略与计算设计协调调度是关键。然而，由于不合理的采样策略，现有方法面临准确性与速度之间的权衡。此外，随着计算优化的推进，数据加载逐渐成为GNN推理中的主要瓶颈。为此，我们提出AES-SpMM，一种自适应边采样SpMM内核。它考虑了每行非零元素数量与共享内存宽度之间的关系。根据每行的不同情况，自适应选择边采样方案。AES-SpMM通过自适应边采样减小图的规模以适配GPU的共享内存，从而降低计算成本并提高数据局部性，从而在GNN推理中实现准确性与速度的平衡。此外，我们引入基于量化的AES-SpMM，将量化和反量化应用于GNN中的特征数据。该方法显著减少了数据加载时间，同时保持准确性损失可接受。我们用常见的GNN模型和数据集评估了AES-SpMM。结果表明，AES-SpMM相对于cuSPARSE SpMM内核和GE-SpMM分别提高了最高可达25.87倍和23.01倍，准确性损失不足1%。相较于ES-SpMM，准确性损失平均降低了3.4%，实现了1.31倍的加速。相较于AES-SpMM，基于量化的AES-SpMM最大准确性损失为0.3%，特征数据加载时间开销降低了50.91%-70.51%。
引用
@article{arxiv.2503.18427,
title = {AES-SpMM: Balancing Accuracy and Speed by Adaptive Edge Sampling Strategy to Accelerate SpMM in GNNs},
author = {Yingchen Song and Yaobin Wang and Yi Luo and Huan Wu and Pingping Tang},
journal= {arXiv preprint arXiv:2503.18427},
year = {2025}
}