Accelerating Physics Simulations with TPUs: An Inundation Modeling Example

Damien Pierce, R. Lily Hu, Yusef Shafi, Anudhyan Boral, Vladimir Anisimov, Sella Nevo, Yi-fan Chen

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.