[논문 리뷰] High-resolution imaging on TPUs

Fantine Huot, Yifan Chen외 3명|arXiv (Cornell University)|2019. 12. 13.

Seismic Imaging and Inversion Techniques참고 문헌 19인용 수 8

한 줄 요약

이 논문은 고차수 수치 스텐실과 TPU의 고대역폭 인터칩 네트워크를 활용하여 클라우드 TPU가 전체파면역진(FWI)을 효율적으로 가속화할 수 있음을 보여준다. 저자들은 텐서플로우를 사용하여 3D 파동장 모델링에서 경쟁력 있는 성능—테슬라 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.

연구 동기 및 목표

• 딥러닝을 초월한 과학적 영상 처리에 클라우드 TPU를 활용할 수 있는지 탐색한다.
• FWI의 높은 계산 비용을 TPU 전용 최적화 기법을 통해 해결한다.
• TPU가 3D 영상의 파동장 모델링에서 GPU 성능을 따라하거나 능가할 수 있는지 평가한다.
• 과학 계산을 위한 이식성과 확장성이 뛰어난 고수준의 읽기 쉬운 텐서플로우 구현체를 개발한다.

제안 방법

• 텐서플로우를 활용해 클라우드 TPU에서 실행 가능한 전체파면역진(FWI) 파동장 모델링 엔진을 개선하였다.
• 정확도 향상과 계산 도메인 크기 감소를 위해 고차수 수치 스텐실(8차 및 16차)을 적용하였다.
• 파동 방정식 해법기에서 밀도 있는 수치 커널을 위한 TPU의 효율적인 행렬 곱셈 기능을 활용하였다.
• 스텐실 계산에서 효율적인 할로 교환을 위해 전용 고속, 저지연 인터칩 메시 네트워크를 활용하였다.
• TPU 활용도와 메모리 대역폭을 극대화하기 위해 텐서 차원을 128의 배수로 최적화하였다.
• 모든 워크플로우를 고수준 파이썬 코드로 구현하여 가독성과 확장성의 용이성을 확보하였다.

실험 결과

연구 질문

• RQ1 클라우드 TPU는 계산 비용이 매우 높은 전체파면역진(FWI)을 효율적으로 가속화할 수 있는가?
• RQ2 고차수 스텐실의 사용은 파동장 모델링에서 TPU 활용도와 계산 도메인 크기에 어떤 영향을 미치는가?
• RQ3 할로 교환에 TPU의 고대역폭 인터칩 네트워크를 활용할 경우 어떤 성능 향상 효과를 얻을 수 있는가?
• RQ4 TPU 기반 FWI 성능은 최신 GPU 구현체와 비교해 속도 및 효율성 면에서 어떻게 비교되는가?
• RQ5 고수준의 텐서플로우 구현체로 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의 성능을 기록하여 강력한 확장성과 효율성을 입증하였다.
• 고수준 파이썬 코드로 구현된 TPU 코드는 고도로 최적화된 C++/CUDA GPU 기준선과 비교해도 경쟁력 있는 성능을 달성하였다.
• 결과적으로 TPU는 고차수 스텐실 및 효율적인 할로 교환 최적화와 결합될 경우 과학 계산 워크로드에 실질적인 대안이 될 수 있음을 시사한다.