[论文解读] FeatGraph: A Flexible and Efficient Backend for Graph Neural Network Systems

Yuwei Hu, Zihao Ye+7 more|arXiv (Cornell University)|Aug 26, 2020

Advanced Graph Neural Networks参考文献 43被引用 12

一句话总结

FeatGraph 是一种高性能图神经网络（GNN）后端，通过两层编程接口将图遍历与特征维度计算联合优化，结合粗粒度稀疏模板（广义 SpMM 和 SDDMM）与细粒度用户自定义函数（UDFs），在 CPU 上实现最高 32× 的加速，在 GPU 上实现最高 7× 的加速，通过同时利用图结构和张量级并行性。

ABSTRACT

Graph neural networks (GNNs) are gaining increasing popularity as a promising approach to machine learning on graphs. Unlike traditional graph workloads where each vertex/edge is associated with a scalar, GNNs attach a feature tensor to each vertex/edge. This additional feature dimension, along with consequently more complex vertex- and edge-wise computations, has enormous implications on locality and parallelism, which existing graph processing systems fail to exploit. This paper proposes FeatGraph to accelerate GNN workloads by co-optimizing graph traversal and feature dimension computation. FeatGraph provides a flexible programming interface to express diverse GNN models by composing coarse-grained sparse templates with fine-grained user-defined functions (UDFs) on each vertex/edge. FeatGraph incorporates optimizations for graph traversal into the sparse templates and allows users to specify optimizations for UDFs with a feature dimension schedule (FDS). FeatGraph speeds up end-to-end GNN training and inference by up to 32x on CPU and 7x on GPU.

研究动机与目标

• 解决由于稀疏操作和特征张量计算处理效率低下导致的 GNN 工作负载性能瓶颈。
• 克服现有深度学习框架的局限性，这些框架在基本稀疏操作之外缺乏对多样化 GNN 内核的灵活支持。
• 弥合传统图处理系统与 GNN 工作负载之间的差距，通过考虑特征张量维度这一根本性改变计算模式的因素。
• 通过将图遍历与特征维度计算的优化解耦并可组合，实现在 CPU 和 GPU 上的高性能 GNN 训练与推理。
• 提供一种可移植、可扩展的后端，可集成到现有 GNN 框架（如 DGL 和 PyTorch Geometric）中。

提出的方法

• 引入两种广义稀疏模板：用于顶点级计算的广义 SpMM 和用于边级计算的广义 SDDMM，支持在特征张量上使用任意用户自定义函数（UDFs）。
• 将内核规范分解为粗粒度的遍历逻辑（通过模板）与细粒度的特征计算（通过 UDFs），实现在每一层独立优化。
• 应用图划分与缓存感知优化，以提升 CPU 上的数据局部性并减少内存访问开销。
• 针对稀疏模式调整并行化策略，以充分利用 GPU 的大规模并行能力，尤其是在特征维度计算方面。
• 采用张量编译器方法，在 CPU 和 GPU 上自动搜索最优调度与代码生成策略。
• 通过暴露清晰、可扩展的接口，支持与现有 GNN 框架的灵活集成，以实现后端加速。

实验结果

研究问题

• RQ1 当 GNN 工作负载同时涉及稀疏图结构与密集特征张量时，如何高效表达与优化？
• RQ2 在现有深度学习与图处理系统中，处理 GNN 特定内核时的关键性能瓶颈是什么？
• RQ3 图遍历与特征维度计算的联合优化能否显著提升端到端 GNN 训练与推理性能？
• RQ4 两层编程抽象（模板 + UDFs）在多大程度上能够为多样化 GNN 模型同时提供灵活性与高性能？
• RQ5 针对特征维度计算与图遍历的针对性优化，在 CPU 和 GPU 平台上如何共同影响性能？

主要发现

• FeatGraph 在端到端 GNN 训练与推理工作负载中，相比最先进系统，在 CPU 上最高实现 32× 加速，在 GPU 上最高实现 7× 加速。
• 广义 SpMM 与 SDDMM 模板结合可自定义 UDFs，能够有效支持多样化 GNN 模型，包括 MLP 聚合与注意力机制。
• 缓存感知的图划分与特征维度感知的负载均衡显著提升了 CPU 性能，通过减少内存访问开销。
• GPU 性能通过同时利用图遍历与特征张量计算中的并行性得到显著增强，而此前的 GPU 图系统大多忽略了这一点。
• FeatGraph 在 GNN 特定内核上优于厂商优化库（如 cuSPARSE 和 MKL），因其能够联合优化遍历与特征计算。
• 该系统具有可移植性，可作为高性能后端集成到现有 GNN 框架（如 DGL 和 PyTorch Geometric）中。