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

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.

• The paper introduces WaterLily.jl, a novel differentiable CFD simulator that integrates implicit LES and machine learning using Julia’s innovative features.
• It employs a Cartesian-grid method with a third-order finite volume scheme and a geometric multigrid solver, achieving GPU speed-ups up to 182x over serial CPU execution.
• The simulator supports heterogeneous execution across serial, multi-threaded CPU, and GPU platforms, paving the way for seamless ML-CFD integration.

WaterLily.jl: A Differentiable Fluid Simulator in Julia

WaterLily.jl presents a novel approach to computational fluid dynamics (CFD), leveraging the Julia programming language to bridge the gap between high-performance computing and flexible, high-level programming required for machine learning integration. This paper focuses on the development of an efficient, open-source simulator for incompressible viscous flow, notable for its support of heterogenous execution across various platforms including serial CPU, multi-threaded CPU, and GPU.

The core innovation of this work lies in its use of Julia, a language designed for scientific computing that combines both high-level usability and low-level performance. By utilizing Julia’s advanced features such as multiple-dispatch, dynamic typing, and meta-programming, WaterLily.jl overcomes the limitations traditionally encountered by CFD software based on rigid low-level languages like C and Fortran.

Computational Strategy and Implementation

WaterLily employs a Cartesian-grid method to solve the Navier-Stokes equations, utilizing a third-order finite volume scheme. This framework includes advanced techniques such as implicit Large Eddy Simulation (LES) for turbulence modeling and the Boundary Data Immersion Method for accurate boundary condition enforcement. The solver’s pressure system is addressed via a Geometric-MultiGrid (GMG) approach, achieving linear scaling with increasing problem size.

A standout feature of WaterLily is its implementation of Automatic Differentiation (AD), allowing the entire solver to be differentiable. This capability facilitates the integration of machine learning models, optimization procedures, and novel data-driven methods directly into the CFD framework.

Performance Results

The adaptability of WaterLily's architecture is evident in its performance benchmarks. The simulator demonstrates impressive speed-ups when executed on GPUs, achieving performance gains up to 182 times compared to serial CPU execution, and up to 23 times with multi-threaded CPU execution. These results indicate that WaterLily not only matches but can exceed the performance of classical Fortran-based solvers in several test cases, including highly challenging scenarios such as the Taylor-Green vortex and flow past a donut geometry.

The performance evaluation underscores WaterLily's effectiveness in handling both computational fluid dynamics and machine learning workloads across different computational backends. This positions the simulator as a potent tool for a range of applications, from fundamental fluid dynamics research to high-level AI-driven optimization tasks.

Implications and Future Work

The development of WaterLily marks a significant step towards more integrative and versatile CFD tools. By combining the computational power of GPU architectures with the flexibility of a high-level language like Julia, this work opens avenues for more seamless integration between CFD and machine learning. Future advancements could explore more extensive use of data-driven techniques for turbulence modeling and optimization, leveraging WaterLily's fully differentiable nature.

Future work may also focus on addressing the limitations of kernel-driven approaches used in WaterLily, particularly the small kernel size which may restrict efficient utilization of computing resources. Exploring more complex, integrated kernel architectures could yield further performance improvements, particularly for large-scale simulations.

In conclusion, WaterLily.jl represents a compelling contribution to the CFD software landscape, poised to enable a new class of applications at the intersection of fluid dynamics and machine learning. Its robust design, utilizing Julia’s unique capabilities, suggests promising developments in both the practical application of CFD in industry and the theoretical exploration of fluid dynamics phenomena.