[論文レビュー] High-resolution imaging on TPUs

Fantine Huot, Yifan Chen+3 more|arXiv (Cornell University)|Dec 13, 2019

Seismic Imaging and Inversion Techniques参考文献 19被引用数 8

ひとこと要約

この論文は、高次数の数値ステンシルとTPUの高帯域幅インタチップネットワークを活用することで、Cloud TPUがフルウェーブフォームインバージョン（FWI）—高分解能イメージング技術—を効率的に高速化できることを示している。著者らはTensorFlowを用いて3D波場モデリングにおいて、Tesla V100 GPUと同等またはそれを上回る性能を達成し、TPUがディープラーニングを超えた計算が重い科学的イメージング用途に対しても実用的であることを示している。

ABSTRACT

The rapid evolution of artificial intelligence (AI) is leading to a new generation of hardware accelerators optimized for deep learning. Some of the designs of these accelerators are general enough to allow their use for other computationally intensive tasks beyond AI. Cloud tensor processing units (TPUs) are one such example. Here, we demonstrate a novel approach using TensorFlow on Cloud TPUs to implement a high-resolution imaging technique called full-waveform inversion. Higher-order numerical stencils leverage the efficient matrix multiplication offered by the Cloud TPU, and the halo exchange benefits from the dedicated high-speed interchip connection. The performance is competitive when compared with Tesla V100 graphics processing units and shows promise for future computation- and memory-intensive imaging applications.

研究の動機と目的

• Cloud TPUをディープラーニングを超える高分解能イメージングに使用する可能性を検討すること。
• FWIの高い計算コストを、TPU固有の最適化を活用して軽減すること。
• 3Dイメージングの波場モデリングにおいて、TPUの性能がGPUを上回るか同等であるかどうかを評価すること。
• 科学的計算のためのTPU上で、ポータブルで拡張可能な高水準のTensorFlow実装を構築すること。

提案手法

• TensorFlowを用いて、Cloud TPU上で動作するフルウェーブフォームインバージョン（FWI）波場モデリングエンジンを適応させた。
• 精度の向上と計算ドメインサイズの削減を目的に、高次数の数値ステンシル（8次および16次）を採用した。
• 波動方程式ソルバーにおける密行列演算のためのTPUの効率的な行列乗算機能を活用した。
• ステンシル計算におけるハローエクスチェンジを効率的に行うために、専用の高速・低遅延のチップ間メッシュネットワークを活用した。
• TPUの利用度とメモリ帯域幅を最大化するために、テンソルの次元を128の倍数に最適化した。
• 可読性と拡張のしやすさを確保するため、すべてのワークフローを高水準のPythonコードで実装した。

実験結果

リサーチクエスチョン

• RQ1 Cloud TPUは、計算が重いイメージング技術としてのフルウェーブフォームインバージョン（FWI）を効率的に高速化できるか？
• RQ2 高次数ステンシルの使用が、波場モデリングにおけるTPUの利用度と計算ドメインサイズに与える影響は何か？
• RQ3 ハローエクスチェンジにTPUの高帯域幅インタチップネットワークを活用することで、どの程度の性能向上が得られるか？
• RQ4 速度と効率の観点から、TPUベースのFWI性能は、最先端のGPU実装と比べてどうなるか？
• RQ5 高水準のTensorFlow実装によって、TPU上で競争力のある性能を達成しつつ、保守可能で拡張可能なコードを実現できるか？

主な発見

• 16次ステンシルを用いることで、8次ステンシルと比較して各軸で計算ドメインを12%削減でき、3Dイメージングワークロードにおいて顕著な節約が可能となった。
• 1024³のドメインにおいて、TPU実装はGPUベースラインの約50%の速度（100ステップで15.68秒対8.98秒）で実行された。
• 16次ステンシルに切り替え、ドメインを768³に縮小することで、TPU実装の計算時間は9.14秒にまで短縮され、8次ステンシルのTPU実装を上回り、GPU性能に近づいた。
• 768³ドメインで16次ステンシルを用いた場合、TPU実装は4.96 Gcells/sを達成し、強力なスケーラビリティと効率性を示した。
• 高水準のPythonで実装されたにもかかわらず、高度に最適化されたC++/CUDA GPUベースラインと比較して、競争力のある性能を達成した。
• 結果から、TPUは特に高次数ステンシルや効率的なハローエクスチェンジといったアルゴリズム的最適化と組み合わせることで、科学的計算ワークロードの代替手段として有効であることが示唆された。