TC-GNN：在 GPU 上桥接稀疏 GNN 计算与稠密张量核
机器学习
2023-06-02 v4 分布式、并行与集群计算
摘要
近来，图神经网络（GNNs）作为基于图的机器学习的骨干，在各个领域（如电子商务）展现出巨大成功。然而，由于高度稀疏和不规则的基于图的操作，GNN 的性能通常不尽如人意。为此，我们提出 TC-GNN，首个基于 GPU 张量核单元（TCUs）的 GNN 加速框架。核心思想是将“稀疏”的 GNN 计算与高性能“稠密”的 TCU 相协调。具体而言，我们对主流 GNN 计算框架中的稀疏操作进行了深入分析。我们引入一种新颖的稀疏图转换技术，以促进 TCU 处理稀疏 GNN 工作负载。我们实现了有效的 CUDA 核与 TCU 协作设计，以充分利用 GPU 资源。我们将 TC-GNN 与 PyTorch 框架集成，以获得高可编程性。严谨的实验表明，在各种模型和数据集上，相较最先进的 DGL 框架平均加速 1.70 倍。
引用
@article{arxiv.2112.02052,
title = {TC-GNN: Bridging Sparse GNN Computation and Dense Tensor Cores on GPUs},
author = {Yuke Wang and Boyuan Feng and Zheng Wang and Guyue Huang and Yufei Ding},
journal= {arXiv preprint arXiv:2112.02052},
year = {2023}
}
备注
Paper is accepted to USENIX ATC'23