[論文レビュー] Accelerating SpMM Kernel with Cache-First Edge Sampling for Graph Neural Networks

Chien‐Yu Lin, Liang Luo+1 more|arXiv (Cornell University)|Apr 21, 2021

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

ひとこと要約

本稿では、GNN推論を高速化するため、グラフをGPU共有メモリに収納できるようにエッジをサンプリングするキャッシュ最優先のエッジサンプリング機構と、それにコードデザインされたSpMMカーネルであるES-SpMMを提案する。この手法は、計算量を削減し、キャッシュ局所性を向上させる。標準的なGNNベンチマークにおいて、cuSPARSEより最大4.35倍の高速化を達成し、精度損失なしで、1%未満の精度損失で45.3倍の高速化を実現した。

ABSTRACT

Graph neural networks (GNNs), an emerging deep learning model class, can extract meaningful representations from highly expressive graph-structured data and are therefore gaining popularity for wider ranges of applications. However, current GNNs suffer from the poor performance of their sparse-dense matrix multiplication (SpMM) operator, even when using powerful GPUs. Our analysis shows that 95% of the inference time could be spent on SpMM when running popular GNN models on NVIDIA's advanced V100 GPU. Such SpMM performance bottleneck hinders GNNs' applicability to large-scale problems or the development of more sophisticated GNN models. To address this inference time bottleneck, we introduce ES-SpMM, a cache-first edge sampling mechanism and codesigned SpMM kernel. ES-SpMM uses edge sampling to downsize the graph to fit into GPU's shared memory. It thus reduces the computation cost and improves SpMM's cache locality. To evaluate ES-SpMM's performance, we integrated it with a popular GNN framework, DGL, and tested it using representative GNN models and datasets. Our results show that ES-SpMM outperforms the highly optimized cuSPARSE SpMM kernel by up to 4.35x with no accuracy loss and by 45.3x with less than a 1% accuracy loss.

研究の動機と目的

• GPU上で遅いスパース・デンス行列積（SpMM）が引き起こすGNN推論の性能ボトルネックを解消すること。
• モデルの精度を損なわず、SpMMの計算時間を短縮し、キャッシュ局所性を向上させること。
• エッジサンプリングをSpMMカーネルに直接統合することで、事前処理のオーバーヘッドを排除すること。
• ユーザーのコードを変更せずに、DGLのような既存のGNNフレームワークにシームレスに統合できること。
• 多様なGNNワークロードに適用可能な汎用的で高性能なSpMMソリューションを提供すること。

提案手法

• CPUの事前処理を排除し、GPU上でサンプリングを実行するカーネル内エッジサンプリング機構を導入し、GPUの並列処理を活用する。
• 共有メモリサイズをサンプリングの上限として設定するキャッシュ最優先設計を採用し、サンプリングされたグラフが完全に共有メモリ内に収まるようにする。
• 効率的な共有メモリへのデータロードを実現するため、連続するグローバルメモリアクセスパターンを採用する。
• データセットの特性に応じて速度と精度のバランスを取る2つの軽量なサンプリング戦略（BucketおよびFastRand）を採用する。
• フォーマット変換のオーバーヘッドを回避するため、標準的なCSRフォーマットをサポートし、DGLへのドロップインリプレースを可能にする。
• サンプリング論理とSpMMカーネルをコードデザインすることで、メモリアクセスと計算削減の両面で一貫した最適化を実現する。

実験結果

リサーチクエスチョン

• RQ1 SpMMカーネルにエッジサンプリングを統合することで、事前処理のオーバーヘッドを排除し、推論を高速化できるか？
• RQ2 共有メモリに収まるように設計されたキャッシュ最優先のサンプリング戦略が、SpMMの性能を顕著に向上できるか？
• RQ3 カーネル内エッジサンプリングを用いることで、最小限の精度損失でSpMMの性能をどの程度向上できるか？
• RQ4 cuSPARSEのような高度に最適化されたSpMMカーネルと比較して、ES-SpMMは速度と精度の両面でどのように差をつけるか？
• RQ5 ES-SpMMは、ユーザーのコードを変更せずに、既存のGNNフレームワークにシームレスに統合できるか？

主な発見

• 標準的なGNNベンチマークにおいて、高度に最適化されたcuSPARSE SpMMカーネル比で最大4.35倍の高速化を達成し、精度損失なしに実現した。
• 1%未満の精度損失で、cuSPARSE比で45.3倍の高速化を達成し、性能と精度のトレードオフの点で優れた結果を示した。
• 標準的なGCNモデルでは、推論時間の大部分（最大95%）がSpMMに費やされており、SpMMが主な性能ボトルネックであることが確認された。
• CPUベースの事前処理と比較して、カーネル内サンプリング機構はオーバーヘッドを低減し、より高速で効率的なサンプリングを可能にした。
• ES-SpMMは、GCNとGraphSAGEの複数のGNNモデルおよびReddit、Cora、PubMed、Ppiの複数のデータセットにおいて、高い性能を維持した。
• DGLへの統合は透明で、ユーザーのコード変更が不要であり、フォーマット変換なしに標準的なCSRフォーマットをサポートした。