[论文解读] 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波场模拟，取得了与NVIDIA 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进行高分辨率成像的可行性。
• 通过利用TPU特有的优化手段，解决全波形反演（FWI）带来的高计算成本问题。
• 评估TPU在3D成像的波场模拟中是否能够达到或超越GPU的性能表现。
• 开发一种高层级、可读性强的TensorFlow实现，确保其在TPU上科学计算的可移植性和可扩展性。

提出的方法

• 使用TensorFlow将全波形反演（FWI）波场模拟引擎适配至Cloud TPU上运行。
• 采用更高阶的数值差分格式（8阶和16阶）以提升精度并减小计算域规模。
• 利用TPU高效的矩阵乘法能力，加速波方程求解器中的密集数值核函数。
• 利用TPU专用的高速、低延迟芯片间网格网络，实现差分计算中边界交换的高效处理。
• 优化张量维度，使其为128的倍数，以最大化TPU利用率和内存带宽。
• 采用高层级Python代码实现完整工作流，确保代码可读性与易于扩展。

实验结果

研究问题

• RQ1 Cloud TPU能否高效加速全波形反演（FWI）这一计算密集型成像技术？
• RQ2 在波场模拟中，高阶差分格式对TPU利用率和计算域规模有何影响？
• RQ3 通过利用TPU的高带宽芯片间网络实现边界交换，能获得多大的性能提升？
• RQ4 与最先进的GPU实现相比，基于TPU的FWI在速度和效率方面表现如何？
• RQ5 在TPU上使用高层级TensorFlow实现是否能在保持可维护性和可扩展性的同时，实现具有竞争力的性能？

主要发现

• 与8阶差分格式相比，16阶差分格式在每个轴向上将计算域缩小了12%，显著降低了3D成像工作负载的计算量。
• 在1024³计算域上，TPU实现的运行速度约为GPU基线的50%（100个时间步长耗时15.68秒 vs. 8.98秒）。
• 通过切换至16阶差分格式并将计算域缩减至768³，TPU实现的计算时间缩短至9.14秒，优于8阶差分格式的TPU实现，并接近GPU性能。
• 在768³计算域上，使用16阶差分格式的TPU实现达到了4.96 Gcells/s的计算速率，展现出强大的可扩展性和效率。
• 尽管采用高层级Python代码实现，TPU代码的性能仍与高度优化的C++/CUDA GPU基线实现相当。
• 结果表明，TPU是科学计算工作负载的可行替代方案，尤其在结合高阶差分格式和高效边界交换等算法优化时更为显著。