专题解读 | 图神经网络GPU内核加速

深度神经网络（DNN）的训练和推理得益于 GPU 的加速。GPU 是功能强大的硬件加速器，GPU 内核是指经过优化可在 GPU 的许多处理单元上并行执行的专用程序，从而使深度学习（DL）的计算更加迅速、高效。而与大多数深度学习任务不同，图神经网络（GNN） 计算遵循消息传递范式，非常需要高效的稀疏操作；然而，现有的深度学习框架（如 PyTorch [1]等）和基础设施无法充分处理这一问题。目前，两个广泛使用的GNN框架，PyG [2]和 DGL [3]，在一定程度上缓解了这些问题，并极大地方便了 GNN 算法的设计与实现。但如何进一步提高 GPU 的利用率，加速图神经网络的训练和推理，仍然具有一定的挑战性。近年来，一些工作提出了了一些新颖的 GPU 内核来加速 GNN 的计算。下表总结了针对不同 GNN 的内核工作。

此次解读将通过对这些工作的介绍来展现 GNN 的 GPU 内核优化方面的工作以及前景。

1 前置知识

1.1 稀疏矩阵表示

由于现实世界的图具有高度稀疏的特点，使用传统的邻接矩阵表示一张图通常会包含大量的“0”，在程序设计中会造成极大的内存浪费。因此，在程序设计中，人们大多会采用几种特殊的表示方法（稀疏矩阵，例如COO、CSR、CSC）来表示一张图。

最简单的稀疏矩阵存储方式，一般会采用三元组 (row, col, data) 来保存邻接矩阵中的非零信息。其中 row 对应邻接矩阵中非零信息的行序号，col 对应邻接矩阵的列序号，data对应邻接矩阵中非零信息的值，举例如下：

按行压缩的稀疏矩阵存储方式，由三个一维数组 rowptr, col, data 组成。这种格式要求矩阵元按行顺序存储，rowptr 对应邻接矩阵每一行非零元素的个数（起始通常为0），举例如下：

按列压缩的稀疏矩阵存储方式，由三个一维数组 row, colptr, data 组成。这种格式要求矩阵元按列顺序存储，colpt 对应邻接矩阵每一列非零元素的个数（起始通常为0），举例如下：

1.2 消息传递范式

假定图的节点集为，边集为，每一个节点有一个特征表示，边有一个特征表示，其中和是特征的维度。消息传递过程如下：

2 以分区为中心的图卷积神经网络加速

2.1 背景问题

受卷积神经网络（CNN）的启发，研究人员将卷积方法应用到图数据结构中，成就了图卷积神经网络（GCN）[4]。GCN 既利用了特征信息又利用了图结构，在很多现实世界的任务中取得了令人信服的结果。然而，由于图数据结构的复杂性，很难在 GPU 上高效地进行 GCN 计算。因此，研究人员提出了一个以分区为中心的图卷积神经网络加速方法（Partition-Centric Graph Convolutional Network, PCGCN）[5]，该方法根据现实世界图通常具有聚类属性的特点，将整个图划分为块，并使用稀疏子图和密集子图的分别处理不同情况的图。

2.2 方法

• 使用局部感知算法（例如 METIS [6]）作为图切分方法来增强局部性
• 对每一个子图应用 GCN 消息传递，生成每一个子图目标节点的嵌入表示
• 将每一个目标节点的所有嵌入表示进行拼接，得到最终表示

详细算法如下：

• SpMM（稀疏矩阵-矩阵乘法）用于稀疏图
• GeMM（通用矩阵-矩阵乘法）用于密集图

• 对于 SpMM 实现方法：

• 首先，CSR 格式的边权重被顺序加载到共享内存中以进行快速访问
• 然后，调度连续的GPU线程来处理顶点特征向量和边权重的乘法，其中连续的线程将处理特征向量的连续维度并更新累积的中间数据寄存器。
• 最后，处理完顶点的所有传入边后，寄存器中累积的中间数据将被写回全局内存。

• 对于 GeMM 实现方法：

  和 SpMM 的区别在于计算子图将不规则的稀疏计算转换为规则的密集计算

• 将一个子图调度到一个 GPU 线程块上进行处理。密集格式的边数据将被加载到共享内存中以进行快速访问。
• 之后，连续的 GPU 线程被调度来处理顶点特征向量和边权重的乘法，并更新寄存器中累积的中间数据，
• 最后合并内存访问并减少全局内存中的写入操作。

2.3 效果

• 与 GTX 1080Ti 上的最佳基准相比，平均加速 2.9 倍（高达 8.8 倍）
• 与 RTX 2080Ti 的最佳基准相比，平均加速 2.9 倍（高达 7.7 倍）

3 内核融合的图卷积神经网络加速

3.1 背景问题

GNN 计算过程可以抽象成三个不同的阶段：消息生成、邻居聚合、特征更新。对于消息生成，通常只是稠密特征矩阵的乘法，因此现有的许多框架已经有了非常好的支持。然而，对于邻居聚合，需要引入大量的 DRAM 存储占用以及数据移动，并且会受到内核启动时间的影响，因此本工作提出 FuseGNN [7]，作为 PyTorch 的扩展，开发两种工作流抽象，为 GNN 提供高度优化的 API 和 GPU 内核。

3.2 方法

假设输入图采用 COO 格式：Gather-ApplyEdge-Scatter (GAS)

假设输入图采用 CSR 格式：Gather-ApplyEdge-Reduce (GAR)

具体来说，GAS 与 COO 格式图一起工作，并处理边级聚合；GAR 处理节点聚合，这对于 CSR 格式特别有用。

• GAS 方法：对于边 ，其中 是源顶点， 是目标顶点， 的特征向量是从输入特征矩阵中选择的。应用边权重后，结果分散到输出特征矩阵的相应行 中。
• GAR 方法：对于每个传入边 ，该过程类似于 GAS 模型。唯一的区别是，ApplyEdge 的输出不是将结果分散到输出特征矩阵，而是在片上缓冲区（寄存器或用户管理的数据缓存）中减少。当所有传入边都减少后，缓冲区的内容被写入 DRAM 中的输出特征矩阵。

3.3 效果

• fuseGNN 将存储占用量减少了几个数量级。特别是，与 PyG 和 DGL 相比，峰值 DRAM 使用量分别减少了 95 倍和 26 倍。

4 GPU 上图神经网络计算的轻量级两级并行范式

4.1 背景问题

尽管最近提出了一些加速 GNN 计算的工作，但它们面临着繁重的预处理、低效原子操作和不必要的内核启动的限制。在此工作中，设计了 TLPGNN[8]，一种用于 GNN 计算的轻量级两级并行范例，即第一级的顶点并行性和第二级的特征并行性，并采用一种新颖的混合动态工作负载分配来解决不平衡的工作负载分配问题。

4.2 方法

• 第一级（即 warp 级），TLPGNN 将每个目标节点分配给一个 warp 线程，以避免由于节点度不均匀而导致分支发散。为了计算每个节点的新嵌入
• 第二级（即线程级）采用特征级并行性，即同一 warp 处理中的线程并行处理特征，同时顺序处理每个边，允许合并内存访问每个源的输入特征节点
• 为了实现平衡的工作负载，开发了动态工作负载分配，一旦释放一个硬件资源，就会分配下一个计算任务

4.3 效果

• TLPGNN 比 GNNAdvisor [9] 实现了 7.7 倍（高达 13.5 倍）加速，比 DGL 实现了 5.6 倍（高达 13.9 倍）加速，比 FeatGraph [10] 实现了 3.3 倍（高达 7.5 倍）加速
• 对于边数和平均度数较大的图数据集，TLPGNN 与 DGL 相比平均实现了 3.7 倍的加速，与 Featgraph 相比平均实现了 4.1 倍的加速
• 对于图卷积运算构造相对简单的三个模型（即 GCN、GIN [11] 和 GraphSage [12]），TLPGNN 为 GCN 实现了 5.8 倍，为 GIN 实现了 4.6 倍，为 GraphSage 实现了 4.7 倍
• 对于最复杂的模型 GAT [13]，TLPGNN 的性能优于其他两个基线 6.5 倍和 2.6 倍

5 总结

GPU 作为深度学习的硬件加速器，在近几年来备受关注。图结构可能非常不规则并且可以以不同的格式呈现，因此一些工作利用双模式 GPU 计算和动态 GPU 工作负载分配来处理不同结构的图。相关实验也证明了这些算法的有效性。

参考资料：

[1] Paszke A, Gross S, Massa F, et al. Pytorch: An imperative style, high-performance deep learning library[J]. Advances in neural information processing systems, 2019, 32.

[2] Fey M, Lenssen J E. Fast graph representation learning with PyTorch Geometric[J]. arXiv preprint arXiv:1903.02428, 2019.

[3] Wang M Y. Deep graph library: Towards efficient and scalable deep learning on graphs[C]//ICLR workshop on representation learning on graphs and manifolds. 2019.

[4] Kipf T N, Welling M. Semi-supervised classification with graph convolutional networks[J]. arXiv preprint arXiv:1609.02907, 2016.

[5] Chao Tian, Lingxiao Ma, Zhi Yang, and Yafei Dai. 2020. Pcgcn: Partition-centric processing for accelerating graph convolutional network. In 2020 IEEE International Parallel and Distributed Processing Symposium (IPDPS). IEEE, 936–945.

[6] G. Karypis and V. Kumar, “A fast and high quality multilevel scheme for partitioning irregular graphs,” SIAM Journal on scientific Computing, vol. 20, no. 1, pp. 359–392, 1998.

[7] Chen Z, Yan M, Zhu M, et al. fuseGNN: Accelerating graph convolutional neural network training on GPGPU[C]//Proceedings of the 39th International Conference on Computer-Aided Design. 2020: 1-9.

[8] Fu Q, Ji Y, Huang H H. TLPGNN: A lightweight two-level parallelism paradigm for graph neural network computation on GPU[C]//Proceedings of the 31st International Symposium on High-Performance Parallel and Distributed Computing. 2022: 122-134.

[9] Wang Y, Feng B, Li G, et al. {GNNAdvisor}: An adaptive and efficient runtime system for {GNN} acceleration on {GPUs}[C]//15th USENIX symposium on operating systems design and implementation (OSDI 21). 2021: 515-531.

[10] Hu Y, Ye Z, Wang M, et al. Featgraph: A flexible and efficient backend for graph neural network systems[C]//SC20: International Conference for High Performance Computing, Networking, Storage and Analysis. IEEE, 2020: 1-13.

[11] Xu K, Hu W, Leskovec J, et al. How powerful are graph neural networks?[J]. arXiv preprint arXiv:1810.00826, 2018.

[12] Hamilton W, Ying Z, Leskovec J. Inductive representation learning on large graphs[J]. Advances in neural information processing systems, 2017, 30.

[13] Veličković P, Cucurull G, Casanova A, et al. Graph attention networks[J]. arXiv preprint arXiv:1710.10903, 2017.