WaterLily.jl: A differentiable fluid simulator in Julia with fast heterogeneous execution

TL;DR

WaterLily.jl is an open-source, differentiable fluid simulator in Julia that offers high performance and flexibility across platforms, enabling advanced optimization and machine learning applications in fluid dynamics.

Contribution

It introduces a multi-platform, differentiable fluid solver in Julia with automatic differentiation and GPU acceleration, bridging CFD with machine learning workflows.

Findings

Linear scaling of computational time with degrees of freedom on CPUs

Up to 182x speed-up with CUDA kernels

Performance comparable to Fortran solvers on research problems

Abstract

Integrating computational fluid dynamics (CFD) software into optimization and machine-learning frameworks is hampered by the rigidity of classic computational languages and the slow performance of more flexible high-level languages. WaterLily.jl is an open-source incompressible viscous flow solver written in the Julia language. The small code base is multi-dimensional, multi-platform and backend-agnostic (serial CPU, multi-threaded, & GPU execution). The simulator is differentiable and uses automatic-differentiation internally to immerse solid geometries and optimize the pressure solver. The computational time per time step scales linearly with the number of degrees of freedom on CPUs, and we see up to a 182x speed-up using CUDA kernels. This leads to comparable performance with Fortran solvers on many research-scale problems opening up exciting possible future applications on the cutting edge of machine-learning research.