[论文解读] Deep Graph Library Optimizations for Intel(R) x86 Architecture

Sasikanth Avancha, Vasimuddin+2 more|arXiv (Cornell University)|Jul 13, 2020

Advanced Graph Neural Networks参考文献 8被引用 5

一句话总结

本论文通过利用内存带宽感知的内核融合、向量化操作和底层优化，对 Intel® x86 CPU 上的 Deep Graph Library (DGL) 聚合原语——特别是 Binary-Reduce 和 Copy-Reduce——进行了优化。作者在七个 GNN 应用中实现了每轮训练时间最高达 13 倍的加速，其中 Binary-Reduce 在 CPU 工作负载上最高提升了 34 倍。

ABSTRACT

The Deep Graph Library (DGL) was designed as a tool to enable structure learning from graphs, by supporting a core abstraction for graphs, including the popular Graph Neural Networks (GNN). DGL contains implementations of all core graph operations for both the CPU and GPU. In this paper, we focus specifically on CPU implementations and present performance analysis, optimizations and results across a set of GNN applications using the latest version of DGL(0.4.3). Across 7 applications, we achieve speed-ups ranging from1 1.5x-13x over the baseline CPU implementations.

研究动机与目标

• 解决 DGL 在 CPU 上进行 GNN 推理和训练时的性能瓶颈问题，其中聚合原语占用了大部分运行时间。
• 识别出 Binary-Reduce 和 Copy-Reduce 原语在实现上的不足，这些实现依赖于序列化操作和冗余的内存复制。
• 优化 DGL CPU 后端中的底层原语，以最大化内存带宽利用率并降低延迟。
• 在 Intel Xeon 处理器上使用真实世界和合成数据集，在多种 GNN 工作负载中展示可测量的性能提升。
• 确保优化后的实现保持数值准确性，并与 PyTorch 后端及 DGL 的高层 API 兼容。

提出的方法

• 使用 PyTorch 性能分析器和应用级性能分析，分析 DGL 的聚合流水线，以隔离 Binary-Reduce 和 Copy-Reduce 操作中的性能瓶颈。
• 通过融合内核重新实现 Binary-Reduce 原语，将逐元素操作和归约合并为单一内存访问模式，以减少延迟并提高缓存效率。
• 通过消除冗余的缓冲区复制并利用向量化指令（如 AVX/AVX512）来优化 Copy-Reduce，以提升特征张量操作的性能。
• 应用内核融合和循环分块技术，以减少稀疏-密集矩阵乘法公式在聚合中产生的内存流量，并改善数据局部性。
• 将优化后的原语集成到 DGL v0.4.3 中，并使用标准 GNN 基准测试验证其正确性和性能。
• 对其他显著影响特定应用运行时间的 PyTorch 原语（如 BatchNorm1d、Embedding）进行性能分析和调优。

实验结果

研究问题

• RQ1 尽管具有高层抽象的优势，为何 DGL 的 CPU 基 GNN 应用仍表现出较差的性能？
• RQ2 在 Intel x86 CPU 上，DGL 的 Binary-Reduce 和 Copy-Reduce 原语的主要性能瓶颈是什么？
• RQ3 低层级优化（如内核融合和向量化）如何改善 GNN 聚合中的内存带宽利用率？
• RQ4 在多种 GNN 工作负载中，优化后的原语能在多大程度上减少每轮训练时间？
• RQ5 是否可以在不牺牲数值准确性或框架兼容性的情况下实现性能提升？

主要发现

• Binary-Reduce 原语在 GNN 应用中占用了大部分运行时间，是 CPU 上的主要性能瓶颈。
• 通过内核融合和向量化优化 Binary-Reduce，在 Intel Xeon 8280 CPU 上实现每轮训练时间最高达 34 倍的加速。
• Copy-Reduce 的优化减少了冗余内存复制，改善了数据局部性，显著提升了消息传递操作的性能。
• 在使用批处理的 GraphSAGE 中，优化后的 DGL 实现了 1.5×–1.7× 的整体加速，其中 Binary-Reduce 的性能提升达 7.2×–10.6×。
• 在 LGNN 应用中，优化后的 BatchNorm1d 和 Embedding 原语分别实现了 13× 和 76× 的加速，共同贡献了 2× 的整体训练加速。
• 所有优化均保持与基线 DGL 实现的数值等价性，确保了所有测试的 GNN 应用中正确性。