Google Research

A TensorFlow Simulation Framework for Scientific Computing of Fluid Flows on Tensor Processing Units

  • John Roberts Anderson
  • Matthias Ihme
  • Qing Wang
  • Yi-fan Chen
Computer Physics Communications (2022)


A computational fluid dynamics (CFD) simulation framework for predicting complex flows is developed on the Tensor Processing Unit (TPU) platform. The TPU architecture is featured with accelerated performance of dense matrix multiplication, large high bandwidth memory, and a fast inter-chip interconnect, which makes it appealing for scientific high performance computing (HPC). This CFD framework is implemented with the finite difference method on a collocated structured mesh, and it uses the graph-based TensorFlow as the programming library. The accuracy and speed of this framework is studied both numerically and analytically. In particular, the impact of machine precision on the mesh convergence is discussed in detail, which provides a guideline for the simulation framework to achieve the desired rate of convergence. Additionally, a linear weak scaling and a super-linear strong scaling is observed up to a full TPU v3 pod with 2048 cores. The algorithm and implementation are validated with canonical 2D and 3D Taylor Green vortex simulations. To demonstrate the capability for simulating turbulent flows, verification simulations are conducted for two configurations, namely the decaying homogeneous isotropic turbulence and a turbulent planar jet. Both simulations show good statistical agreement with reference solutions.

Research Areas

Learn more about how we do research

We maintain a portfolio of research projects, providing individuals and teams the freedom to emphasize specific types of work