[論文レビュー] FusionStitching: Deep Fusion and Code Generation for Tensorflow Computations on GPUs
FusionStitching は、ブロック構成と共有メモリ最適化を用いて効率的な GPU カーネル統合を実現する Tensorflow 専用のディープフェージュニオンおよびコード生成システムであり、XLA と比較して GPU カーネル起動を 55%(幾何平均)削減し、レイテンシが重要なワークロードにおいてエンドツーエンド性能を最大 20% 向上させます。
In recent years, there is a surge on machine learning applications in industry. Many of them are based on popular AI frameworks like Tensorflow, Torch, Caffe, or MxNet, etc, and are enpowered by accelerator platforms such as GPUs. One important challenge of running Tensorflow computations on GPUs is the fine granularity problem, namely, FLOPS of individual ops are far from enough to fully exploit the computing power of underlying accelerators. The XLA framework provides a solid foundation to explore this problem further. In this paper, we propose FusionStitching, a novel, comprehensive Op fusion and code generation system to stitch computations into large GPU kernels. Experimental results on four public models and two of our large inhouse applications show another 55% (geometric mean) reduction of GPU kernel launches, compared to the XLA fusion baseline. This increases the E2E performance of both of our latency critical inhouse applications up to 20%.
研究の動機と目的
- GPU ハードウェアを十分に活用できないほど FLOPS が低い個々の演算が生じる TensorFlow の細粒度オペレーション問題に対処すること。
- XLA の静的統合ルールでは処理できない複雑なケース(高コストな要素演算、バッチ化行列積、レイアウト変換など)を克服すること。
- 計算仕様と低レベルのスケジューリング詳細を分離することで、非 MatMul/Convolution オペレーションに対してもスケーラブルかつ自動化された統合とコード生成を可能にすること。
- ブロック構成と共有メモリを活用した新しいコード生成戦略を考案し、異なる並列ループエミッタを備えた複数のオペレーションが統合されたカーネルを構成すること。
- カーネル起動オーバーヘッドを最小限に抑えることで、特にレイテンシに敏感な産業用途において顕著なパフォーマンス向上を達成すること。
提案手法
- データ依存性チェーンの最長パスを特定するための重要なパス解析を提案し、統合意思決定を最長のデータ依存性チェーン上にある演算に優先的に導く。
- 低レベルチューニングパラメータを抽象化しつつ自動最適化を可能にするドメイン特化スケジュール記述言語を導入。
- 複数のオペレーションが別個の並列ループエミッタを使用しつつ、共有メモリを中間手段として利用できるようにする、新しい IR エミッタ IrEmitterStitched を設計。
- 異なるスケジューリングパターンを持つ演算間での効率的なデータ転送を可能にするために、オンチップ共有メモリを一時領域として活用。
- 使用量が 20KB のしきい値を超えた場合に、動的にメモリフットプリントを削減する共有メモリの縮小メカニズムを実装。
- 依存関係解析にインspiredされたレイヤードスパングラフを用いて、複雑な計算グラフ全体における統合の機会をモデル化・最適化。
実験結果
リサーチクエスチョン
- RQ1標準的な統合ヒューリスティクスが複雑な依存関係やメモリアクセスパターンのため失敗する状況下で、GPU 上の細粒度 TensorFlow オペレーションを効果的に統合する方法は何か?
- RQ2共有メモリの構成とブロックレベルのスケジューリングは、ディープラーニングワークロードにおけるカーネル統合のスケーラビリティとパフォーマンスにどのような影響を与えるか?
- RQ3コンパイラー駆動型でユーザーに透明な統合およびコード生成システムは、XLA に基づく従来の統合戦略に比べ、カーネル起動削減とエンドツーエンドレイテンシの面で優れているか?
- RQ4複数の統合オペレーションにわたって共有メモリを再利用することで、メモリ帯域幅の圧迫をどれほど軽減でき、スループットを向上させられるか?
- RQ5提案された統合戦略は、公開モデルおよび大規模産業応用を含む多様なワークロードに、どの程度スケーラブルに拡張可能か?
主な発見
- FusionStitching は、4つの公開モデルと2つの社内アプリケーションにおいて、XLA の統合ベースラインと比較して GPU カーネル起動を幾何平均で 55% 削減した。
- レイテンシが重要な2つの産業用途において、エンドツーエンドパフォーマンスが最大 20% 向上し、実世界での影響を実証した。
- システムは、XLA の静的統合ルールでは問題となる複雑な統合ケース(バッチ化行列積、高コストな要素演算、メモリレイアウト変換など)を正常に処理できた。
- 音声処理ワークロードでは共有メモリ使用量が最大 10,432 バイトに達し、NMT モデルでは 17% の割合が複数のオペレーションで共有されたことで、結果の再利用が効果的に行われた。
- 100 以上のカーネルのうち 3 つのカーネルで縮小メカニズムが発動したが、これはほとんどのワークロードにおいてメモリ圧力が 20KB の制限内に収容可能であることを示している。
- LR、W2V、RNN、BiRNN といったワークロードでは共有メモリ使用量が最小限で、縮小処理も発動しなかったため、単純な統合パターンと安定したパフォーマンスを示した。
より良い研究を、今すぐ始めましょう
論文の読解から最終レビューまで、研究時間を劇的に削減しましょう。
クレジットカード登録不要
このレビューはAIが作成し、人間の編集者が確認しました。