利用 Tensor Core 优化加速稀疏图神经网络
机器学习
2025-02-25 v2 硬件体系结构
摘要
图神经网络（GNNs）在社交网络、生物信息学和推荐系统等领域得到了广泛应用。然而，图数据的不规则性和稀疏性对传统计算方法提出了挑战，这些方法无法满足 GNNs 的性能需求。最近的研究探索了使用 CUDA Cores 和 Tensor Cores 进行并行加速，但仍存在重大挑战：（1）算子融合导致虚假的高利用率，未能将 CUDA Cores 和 Tensor Cores 视为独立资源；（2）异构核心具有不同的计算偏好，导致效率低下。为了解决这些问题，本文提出了 FTC-GNN，这是一个高效利用 CUDA Cores 和 Tensor Cores 进行 GNN 计算的新型加速框架。FTC-GNN 引入了（1）一个协同设计，支持 CUDA Cores 和 Tensor Cores 的并行利用；以及（2）一个稀疏到稠密的转换策略，将稠密矩阵运算分配给 Tensor Cores，同时利用 CUDA Cores 进行数据管理和稀疏边处理。该设计优化了 GPU 资源利用率并提高了计算效率。实验结果证明了 FTC-GNN 在不同数据集上使用 GCN 和 AGNN 模型的有效性。对于 GCN，FTC-GNN 相比于 DGL、PyG 和 TC-GNN 分别实现了 4.90 倍、7.10 倍和 1.17 倍的加速。对于 AGNN，它分别实现了 5.32 倍、2.92 倍和 1.02 倍的加速，确立了其在加速 GNN 计算方面的优越性。
关键词
引用
@article{arxiv.2412.12218,
title = {Accelerating Sparse Graph Neural Networks with Tensor Core Optimization},
author = {Ka Wai Wu},
journal= {arXiv preprint arXiv:2412.12218},
year = {2025}
}