EuroSys（European Conference on Computer Systems）是计算机系统领域的国际顶级学术会议，主要关注计算机系统领域的研究与实践，涉及操作系统、分布式系统等领域。自2006年创办以来，EuroSys 接收并发布了许多高质量的研究论文，代表了该领域内最新的研究进展和创新思想，在国际学术界与工业界形成了广泛影响。
蚂蚁 TuGraph 图学习团队的论文《GPC: 基于编译器的面向图神经网络稀疏计算优化》已被 EuroSys 2023 Poster 接收。
-概述-
在非欧几里得图数据上的图神经网络（GNN）学习最近在许多领域引起了迅速的兴趣。广义稀疏计算对于最大限度地提高GNN学习的性能至关重要，而最近的GNN主要集中在优化与节点、边和额外特征维度相关的粗粒度并行性。然而，有效地实现广义稀疏计算是具有挑战性的。缺乏深入的架构感知设计的广义稀疏计算的性能优化很少被现有的特定领域语言（DSL）所支持，而且因为涉及到大量的试验和错误，很难由专家进行调整。
在这项工作中，我们提出了GPC (General sParse Compiler)，一个新的编译器框架，它扩展了流行的Halide编译器，通过编译器驱动的优化和自动调整，使GNN的广义稀疏计算得到有效加速。为了加速广义稀疏计算，GPC将算法与调度分开，并为现代GPU引入了几种新的稀疏计算优化技术，包括二维共享内存优化、寄存器平铺优化和高效成本驱动的设计空间探索和自动调整。针对高度优化的最先进的稀疏计算内核以及端到端GNN训练和推理效率的广泛评估表明，我们提出的GPC比最先进的稀疏计算内核实现了4.37倍的速度提升，比三个流行的GNN模型GCN、GraphSAGE和GAT的训练和推理速度提升了1.37~2.25倍。
-背景-
图神经网络（GNN）是深度学习中的一个新兴领域，已经应用于广泛的真实世界应用，包括但不限于生物信息学、社会科学、推荐系统和量子化学。使用GNN模型进行图结构表示学习对许多实际的与图相关的应用非常重要，例如节点分类、图同构和扩散模式。与依赖于密集矩阵运算的传统深度学习模型（如DNNs）不同，GNN涉及具有消息感知的传播操作，可以根据其邻居递归地更新每个顶点的特征。在处理真实世界的图数据集时，它们的稀疏性和不规则性对于最近的最先进框架（如TensorFlow、PyTorch和MXNet）构成了性能和设计挑战。
图1. 这是一个在GNN中使用聚合求和的g-SpMM操作示例。左侧是将顶点𝑛1、𝑛2和𝑛3的特征向量聚合到顶点𝑛4上。右侧，我们将图表示为邻接矩阵，以便计算顶点的特征向量作为一个SpMM操作。
通常，GNN的计算模式可以用消息传递原语来描述。例如，DGL引入了广义稀疏计算，具体包括广义采样稠密-稠密矩阵乘法（g-SDDMM）和广义稀疏矩阵-矩阵乘法（g-SpMM），以表达这些消息传递原语。例如，当进行聚合求和操作时，从相邻顶点聚合目标顶点特征可以表示为SpMM操作，如图1所示。类似地，对于目标边的特征聚合可以表示为源顶点特征和目标顶点特征之间的SDDMM操作。这些操作被认为是广义的，因为GNN中的聚合类型可以包括求和、最大值等。然而，这些操作会导致GNN训练带来很大的开销。SpMM和SDDMM样式的操作占总执行时间的60%以上，成为提高GNN学习速度的瓶颈，促使我们加速这些操作。
挑战：优化SpMM和SDDMM样式的运算存在三个挑战，包括当前稀疏矩阵计算优化、计算库和领域特定语言的限制。首先，与采用易于使用的张量驱动并行性的传统密集矩阵计算相比，稀疏矩阵操作由于其不规则的工作负载、稀疏的内存访问和有限的数据重用更具挑战性。其次，在cuSPARSE和Sputnik计算库中虽然考虑了负载平衡和其他特定于稀疏矩阵的优化，但它们主要针对科学计算和常规深度神经网络（DNNs）（如Transformers和ResNet-50），并为GNN中的运算符提供非常有限的支持。最后，现有的领域特定语言（DSLs）例如TVM、Halide和Ansor将算法与调度分离，并利用自动调整策略进行高效的自动代码生成，但对于GNNs并不高效。这些DSL主要针对图像处理和自然语言处理领域，通常使用区间分析，并提供支持处理规则密集矩阵的功能。然而，GNN DSL领域中稀疏计算的优化仍然非常有限。
为了解决这些技术挑战，我们提出了GPC，这是一种基于编译器的方法，可以在现代GPU上高效加速GNN中的广义稀疏计算。GPC是一种新的DSL，扩展了Halide，并能够为通用稀疏计算高效自动生成高度优化的代码。几秒钟内生成的算子性能可以超过专家手动优化数周的代码。总之，我们做出了以下贡献：
• 我们为GNN中的g-SpMM操作提出了2-D共享内存平铺、1-D寄存器平铺和行负载平衡优化。
• 我们为GNN中的g-SDDMM操作提出了自适应warp shuffle优化技术。
• 我们将上述优化策略集成到Halide中，这是一种用于图像处理任务的DSL编译器。
• 我们开发了一种新颖的基于DNN的成本模型来预测性能，并将其与两种自动调整方法（即遗传搜索和随机采样）相结合，用于自动为各种GNN模型生成高度优化的稀疏计算内核。
-实现方法-
图2. GPC 编译流水线。
图2展示了GPC DSL编译器的工作流程。与常规的GNN系统优化不同，我们利用编译器技术来增强g-SpMM和g-SDDMM操作的性能并实现自动调整。这包括优化GPU内存访问、工作负载平衡和SIMD感知执行。我们使用扩展的GPC DSL描述了类似于SpMM和SDDMM的算法以及相应的调度，适用于CUDA-enabled GPU。接下来，我们将使用我们扩展的GPC DSL讨论g-SpMM和g-SDDMM优化及其相应的调度。
图3. GPC对g-SpMM算子的优化，包含Shared Memory Tiling、Register Tiling以及Row balance等优化。
如图3所示，对于g-SpMM的运算，我们引入了缓冲区边界推断技术来实现2-D共享内存优化，引入缓冲区绑定索引表达式来实现负载平衡优化，并使用1-D步幅寄存器平铺来优化寄存器级别的数据重用。在g-SpMM中，表示相邻稀疏矩阵的格式是CSR格式。首先，缓冲区边界推断技术将稀疏矩阵预加载到共享内存中，并重用数据以减少全局内存访问。其次，缓冲区绑定索引表达式绑定图的节点，使得GPC编译器可以将GPU块调度到具有特定数量邻居的特定节点，以优化GNN中的稀疏矩阵计算的平衡问题。最后，1-D步幅寄存器平铺在寄存器中计算多个输出，以便从稀疏矩阵中重用寄存器中的数据。正如表1所示，这些优化是现有计算DSL所未探索的，对于增强稀疏计算的性能至关重要。
表1. 对于g-SpMM操作，2-D共享内存平铺和行负载平衡优化对性能影响最大。对于g-SDDMM操作，是否使用warp shuffle指令对性能至关重要。
图4. g-SDDMM算子以及对应GPC对其的优化。在𝑠𝑡𝑎𝑔𝑒1中，线程𝑡1和𝑡2分别计算矩阵𝑈和𝑉的部分点积结果。𝑆𝑡𝑎𝑔𝑒2使用warp shuffle将𝑡1和𝑡2的部分结果相加以得到最终结果𝑟。
如图4所示，在g-SDDMM的运算中，我们使用COO格式，并应用自适应warp shuffle优化。Warp shuffle函数使用寄存器而不是共享内存或全局内存进行warp内的线程通信。我们将特定数量的边分配给块，使用多个线程将特征计算到中间结果中，并应用自适应warp shuffle将中间结果缩小到寄存器中。
图5. GPC成本模型的DNN架构。
我们还引入了一种基于广泛搜索空间的成本模型自动调整方法，以自动搜索最佳结果。我们的成本模型是一个简单易训练的DNN，如图5所示。模型结构的第一层是对数函数，第二层是批量归一化，接着是三个密集层，最后一层是归一化层。使用均方误差计算损失。真实数据进行归一化以更好地反映真实趋势，使得𝑙𝑜𝑠𝑠 = 𝑀𝑆𝐸 (𝑦𝑝𝑟𝑒𝑑𝑖𝑐𝑡𝑒𝑑 , 𝑛𝑜𝑟𝑚𝑎𝑙𝑖𝑧𝑒 (𝑦))。使用表示并行性的一组度量值，以及GPU块平铺大小、共享内存平铺大小、寄存器平铺大小和自适应warp shuffle的adaptive-lane宽度等，来形成搜索空间。
-实验-
我们使用具有16GB设备内存的NVIDIA V100 GPU，并使用CUDA 11.1版本。邻接稀疏矩阵的索引格式使用32位整型，密集张量的值则使用32位浮点值，以行主格式表示。g-SpMM和g-SDDMM算子使用Hamilton等人的REDDIT数据集以及Open Graph Benchmarks的OGBN-PROTEINS和OGBN-PRODUCTS数据集进行了测试。这三个数据集都被广泛用作评估GNN模型的基准。
我们通过与最先进的GPU实现进行比较来评估GPC，包括DGL v0.9.1、FeatGraph、Sputnik和cuSPARSE，以展示GPC在核函数性能基准、端到端模型训练和推理基准中的性能优势。
图6. REDDIT、OGBN-PROTEINS和OGBN-PRODUCTS数据集的g-SpMM和g-SDDMM基准测试结果。
如图6所示，我们的性能在大多数经过测试的数据集和特征长度上，超过了cuSPARSE、DGL、FeatGraph和Sputnik。
针对g-SpMM运算，如图6 (a) - (c)所示，与使用cuSPARSE实现的DGL相比，我们的实现在REDDIT、OGBN-PROTEINS和OGBN-PRODUCTS上分别实现了平均加速比为3.2×、2.6×和1.5×。而与FeatGraph相比，我们实现了最高3.4×的加速和平均1.9×的加速。与Sputnik相比，我们获得了最高2.1×的加速和平均1.4×的加速。
针对g-SDMM运算, 如图6 (d) - (f)所示，GPC在所有测试情况下的性能都能够始终优于或与DGL和FeatGraph相匹配。GPC在DGL上实现了约1.02×到3.37×的加速，平均加速比为1.46×。
表7. 在REDDIT数据集上，端到端训练（一个epoch）和推理加速，其中GraphSage使用了Full Batch。
另外，我们对比了DGL和GPC在GCN、GAT和GraphSage模型上的端到端性能，数据集使用了Reddit，如表格7所示。可以看到，使用GPC的DGL在所有模型上性能优于原始DGL，训练的平均加速比为1.69×，推理的平均加速比为1.78×。
-总结-
我们提出了GPC来加速GNN中的广义稀疏计算，利用DSL编译器来分离算法和调度。GPC通过引入非矩形缓冲区边界推断和缓冲区绑定索引，扩展了Halide，以实现对稀疏内核的DSL描述和代码生成。此外，GPC在调度中引入了2D共享内存平铺、行负载平衡、1D-stride寄存器平铺、自适应warp shuffle等优化，以提高数据复用效率和显著改善行负载平衡。GPC引入了一种基于DNN的成本模型，结合遗传搜索自动调整，可以在不需要人工干预的情况下自动搜索最优结果。
因此，我们的内核比以前的技术快了多达4.75×。通过将GPC集成到DGL中，与DGL相比，我们的模型训练和推理性能提高了1.37× ∼ 2.25×。我们的自动调整仍需要几秒钟来生成最优程序，并且仅在GPU硬件上使用。利用自动调整生成在各种硬件、数据集和GNN模型上具有最佳性能的程序是未来研究方向之一。
欢迎关注 TuGraph 代码仓库
https://github.com/tugraph-family/tugraph-db
https://github.com/tugraph-family/tugraph-analytics
END
往期回顾
→ 蚂蚁集团开源图数据库TuGraph，成立图计算开源委员会
▼ 关注蚂蚁图计算技术，了解最新资讯
2487
被折叠的 条评论
为什么被折叠?
到【灌水乐园】发言
