[论文解读] TC-GNN: Bridging Sparse GNN Computation and Dense Tensor Cores on GPUs

Yuke Wang, Boyuan Feng+3 more|arXiv (Cornell University)|Dec 3, 2021

Advanced Graph Neural Networks被引用 7

一句话总结

TC-GNN 是首个将稀疏图神经网络（GNN）计算与密集 GPU Tensor Core Units（TCUs）相连接的 GNN 加速框架。它提出了一种新颖的稀疏图转换（SGT）技术，将不规则的稀疏邻居数据压缩为与 TCUs 兼容的密集块，同时联合优化 CUDA 核心与 TCUs，实现高性能的 GNN 推理与训练，在多种模型和数据集上相较最先进的框架 DGL 平均实现 1.70× 的加速。

ABSTRACT

Recently, graph neural networks (GNNs), as the backbone of graph-based machine learning, demonstrate great success in various domains (e.g., e-commerce). However, the performance of GNNs is usually unsatisfactory due to the highly sparse and irregular graph-based operations. To this end, we propose TC-GNN, the first GNN acceleration framework based on GPU Tensor Core Units (TCUs). The core idea is to reconcile the "Sparse" GNN computation with the high-performance "Dense" TCUs. Specifically, we conduct an in-depth analysis of the sparse operations in mainstream GNN computing frameworks. We introduce a novel sparse graph translation technique to facilitate TCU processing of the sparse GNN workload. We implement an effective CUDA core and TCU collaboration design to fully utilize GPU resources. We integrate TC-GNN with the PyTorch framework for high programmability. Rigorous experiments show an average of 1.70X speedup over the state-of-the-art DGL framework across various models and datasets.

研究动机与目标

• 为解决由高度稀疏且不规则的图操作导致的 GNN 性能瓶颈，这些操作在训练和推理中占运行时间的 80% 以上。
• 实现 GPU Tensor Core Units（TCUs）的有效利用，这些 TCUs 专为密集矩阵运算而优化，但在现有稀疏 GNN 框架中利用率较低。
• 通过设计一种能够调和稀疏 GNN 计算与密集 TCU 执行的系统，克服 cuSPARSE 和 cuBLAS 的局限性。
• 提供一种高可编程性的解决方案，与 PyTorch 集成，最大限度减少用户学习成本，同时最大化硬件利用率。
• 探索硬件感知优化，为未来支持动态 TCU 输入形状和结构稀疏性的 AI 加速器提供启发。

提出的方法

• 提出一种新颖的稀疏图转换（SGT）技术，通过利用节点间频繁共享邻居的特性，将稀疏邻居数据重新组织为密集块，减少冗余内存访问。
• 设计一种混合 CUDA 核心与 TCU 内核架构：CUDA 核心负责细粒度内存访问与数据移动，而 TCU 加速对压缩块的密集 GEMM 操作。
• 采用内存与数据流优化，最小化空闲时间并最大化数据重用，尤其适用于高维节点嵌入。
• 通过自定义 API 实现与 PyTorch 的无缝集成，使用户无需更改现有编程工作流即可利用 TC-GNN。
• 通过编译时参数支持 TCU 块大小与精度格式的动态配置，确保在不同 GPU 架构间的适应性。
• 利用近期 CUDA 功能支持使用 TCU 的分块 SpMM，但通过 SGT 预处理将其扩展至处理不规则稀疏图。

实验结果

研究问题

• RQ1 尽管 TCU 专为密集矩阵运算而设计，能否有效用于稀疏 GNN 计算？
• RQ2 如何将现实世界图中固有的不规则性与稀疏性转化为与固定大小 TCU 输入块兼容的形式？
• RQ3 何种系统级设计能够实现 CUDA 核心与 TCU 在 GNN 工作负载中的高效协作，避免因简单应用 TCU 导致的性能下降？
• RQ4 TC-GNN 框架在多种模型与数据集上相较现有最先进 GNN 框架（如 DGL）的加速程度如何？
• RQ5 未来 GPU 硬件应如何增强，以原生支持动态 TCU 块形状与结构稀疏性，从而实现更大的性能提升？

主要发现

• TC-GNN 在多个 GNN 模型与真实世界数据集上相较最先进的 DGL 框架平均实现 1.70× 的加速，证明了显著的性能提升。
• TC-GNN 的吞吐量随节点嵌入维度的增加而成比例提升，表明 GPU 资源得到有效利用，且内存带宽扩展能力强。
• 稀疏图转换（SGT）技术成功通过将共享邻居压缩为更少的密集块，减少了冗余内存访问，提升了数据局部性并降低了开销。
• 混合 CUDA 核心与 TCU 内核设计有效平衡了内存密集型操作（由 CUDA 核心处理）与计算密集型 GEMM（由 TCU 处理），最大化硬件利用率。
• 通过原生 PyTorch 集成，TC-GNN 保持了高可编程性，支持端到端 GNN 工作负载，仅需极少代码修改，显著提升开发效率。
• 该框架可在现代 GPU（如 A100、H100、RTX4090）上跨平台运行，并支持 TCU 参数的可配置性，确保对下一代 GPU 的适应性。