Accelerating Physics Simulations with TPUs: An Inundation Modeling Example

TL;DR

This paper demonstrates that TPUs significantly accelerate physics simulations, specifically flood modeling, achieving over 100 times speedup compared to CPUs, and provides accessible tools via Google Cloud Platform.

Contribution

It introduces the use of TPUs for scientific PDE-based simulations, showcasing substantial speed improvements and providing publicly available interactive notebooks.

Findings

TPUs achieve over 100x speedup over CPUs for flood modeling.

Physics simulations on TPUs are accessible via Google Cloud Platform.

The study extends TPU applications beyond machine learning to scientific computing.

Abstract

Recent advancements in hardware accelerators such as Tensor Processing Units (TPUs) speed up computation time relative to Central Processing Units (CPUs) not only for machine learning but, as demonstrated here, also for scientific modeling and computer simulations. To study TPU hardware for distributed scientific computing, we solve partial differential equations (PDEs) for the physics simulation of fluids to model riverine floods. We demonstrate that TPUs achieve a two orders of magnitude speedup over CPUs. Running physics simulations on TPUs is publicly accessible via the Google Cloud Platform, and we release a Python interactive notebook version of the simulation.