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

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.

• General Science