ECCOMAS 2024

A Flexible Massively Parallel Library for Differentiable Lattice Boltzmann

  • Ataei, Mohammadmehdi (Autodesk)
  • Hennigh, Oliver (NVIDIA)
  • Meneghin, Massimiliano (Autodesk)
  • Salehipour, Hesam (Autodesk)

Please login to view abstract download link

High-performance computing (HPC) is rapidly evolving, not only in capacity but also in complexity due to the emergence of diverse and specialized computing architectures. However, unlocking the full capabilities of these systems for fluid flow simulations has been a complex task, often requiring specialized knowledge. Taking a cue from the machine learning (ML) community's success, we propose a paradigm shift in CFD application development: employing Python's high-level, user-friendly programming interface, while delegating low-level optimizations to robust backends or just-in-time compilers—a strategy that we believe has significantly helped with the recent rapid progress in ML. Introducing the XLB, a CFD library based on the Lattice Boltzmann Method (LBM) that blends the simplicity of Python with the raw power of diverse computing systems. XLB is versatile, ready to work with standard CPUs, GPUs, and even Google's specialized TPUs, as well as advanced platforms like NVIDIA's Grace Hopper. It is built to manage large-scale simulations that stretch across distributed GPU clusters, and it is equipped with technologies such as out-of-core computing to handle the largest of simulation tasks. XLB supports various computational backends and is designed to simplify the integration of new ones, enhancing its adaptability to future technological advancements. The LBM kernels of XLB are fully differentiable, and readily compatible with machine learning libraries like JAX for integrating data-driven models and neural network models directly into fluid dynamics simulations. This differentiability opens the avenues of research for applications requiring optimization, data assimilation, and surrogate modeling. XLB is readily available to the entire scientific community under the permissive Apache 2.0 license.