[論文レビュー] 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 —を最適化している。著者らは、7つの GNN アプリケーションにおいて、1エポックあたりの学習時間を最大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 操作における性能ボトルネックを特定する。
• 要素ごとの演算と削減を1つのメモリアクセスパターンに統合する融合カーネルを用いて、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 ワークロードにおいて、1エポックあたりの学習時間をどの程度短縮できるか？
• RQ5 数値的正確性やフレームワークの互換性を損なわずに、性能向上を達成できるか？

主な発見

• Binary-Reduce プリミティブは GNN アプリケーションにおける実行時間の大部分を占めており、CPU 上での主な性能ボトルネックである。
• カーネル結合とベクトル化による Binary-Reduce の最適化により、Intel Xeon 8280 CPU では 1エポックあたりの時間に最大34倍の高速化が達成された。
• Copy-Reduce の最適化により、不要なメモリコピーが削減され、データ局所性が向上し、メッセージパッシング操作における顕著な性能向上が得られた。
• バッチ処理を伴う GraphSAGE では、最適化された DGL が全体で 1.5×–1.7× の高速化を達成し、Binary-Reduce は 7.2×–10.6× の改善を示した。
• LGNN アプリケーションでは、最適化された BatchNorm1d と Embedding プリミティブがそれぞれ 13× および 76× の高速化を達成し、全体の学習速度が 2倍に向上した。
• すべての最適化は、ベースライン DGL 実装と数値的に同等であることを維持しており、テストされたすべての GNN アプリケーションで正しさが保証された。