[論文レビュー] AsyncTaichi: Whole-Program Optimizations for Megakernel Sparse Computation and Differentiable Programming.

Yuanming Hu, Mingkuan Xu+2 more|arXiv (Cornell University)|Dec 15, 2020

Parallel Computing and Optimization Techniques参考文献 16被引用数 1

ひとこと要約

本稿では、スパースかつ微分可能計算における不要なリスト生成を削減する、Taichi向けの全プログラム最適化フレームワークであるAsyncTaichiを提示する。実行時における計算グラフ全体の解析と複数カーネルに跨る最適化により、コード変更なしにスパースグリッドシミュレーションおよび微分可能プログラミングワークロードで、カーネル起動回数が3.07–3.90倍減少し、1.73–2.76倍の高速化を達成した。

ABSTRACT

We present a whole-program optimization framework for the Taichi programming language. As an imperative language tailored for sparse and differentiable computation, Taichi's unique computational patterns lead to attractive optimization opportunities that do not present in other compiler or runtime systems. For example, to support iteration over sparse voxel grids, excessive list generation tasks are often inserted. By analyzing sparse computation programs at a higher level, our optimizer is able to remove the majority of unnecessary list generation tasks. To provide maximum programming flexibility, our optimization system conducts on-the-fly optimization of the whole computational graph consisting of Taichi kernels. The optimized Taichi kernels are then just-in-time compiled in parallel, and dispatched to parallel devices such as multithreaded CPU and massively parallel GPUs. Without any code modification on Taichi programs, our new system leads to $3.07 - 3.90 imes$ fewer kernel launches and $1.73 - 2.76 imes$ speed up on our benchmarks including sparse-grid physical simulation and differentiable programming.

研究の動機と目的

• Taichiにおけるスパースおよび微分可能プログラミングのパフォーマンスボトル neck を解消するため、不要なリスト生成の原因となる要因を特定する。
• コード変更なしに複数のTaichiカーネルに跨る全プログラム最適化を可能にする。
• カーネル起動のオーバーヘッドを低減し、CPU や GPU などの並列デバイスでの実行効率を向上させる。
• 実行時における最適化を通じて、柔軟かつ高性能なスパース計算ワークロードの実行を実現する。

提案手法

• 高水準でのTaichiプログラムの計算グラフ全体を分析し、余分なリスト生成タスクを同定する。
• 実行時における実行時最適化を実行し、全プログラム最適化を実現する。
• 即時コンパイルにより、マルチスレッドCPUおよびGPU上で並列に実行可能な最適化済みカーネルを生成する。
• 複数のカーネルに跨るスパースグリッドの反復パターンを分析し、余分なリスト生成を排除する。
• カーネル起動回数を最小限に抑え、データ局所性を向上させるために実行グラフを動的に最適化する。

実験結果

リサーチクエスチョン

• RQ1 全プログラム最適化により、スパースTaichiプログラムにおける余分なリスト生成をどのように削減できるか？
• RQ2 コード変更なしに複数カーネルに跨る最適化により、どのようなパフォーマンス向上が達成できるか？
• RQ3 実行時最適化により、スパースおよび微分可能ワークロードにおけるカーネル起動オーバーヘッドをどの程度低減できるか？
• RQ4 物理シミュレーションや微分可能プログラミングを含む多様なワークロードにおいて、最適化フレームワークの有効性はどの程度か？

主な発見

• ベンチマークワークロード全体でカーネル起動回数が3.07～3.90倍削減された。
• スパースグリッド物理シミュレーションおよび微分可能プログラミングタスクで1.73～2.76倍の高速化が達成された。
• 最適化は透明に適用され、既存のTaichiプログラムに変更を加える必要がなかった。
• 本システムは、マルチスレッドCPUおよび大規模並列GPUの両方で効率的な実行を可能にした。