CRUX: A topology-aware load balancer for mesh-based fluid dynamics codes on GPU clusters

The rapid growth of computational power has revolutionized numerical simulations, profoundly enhancing our understanding of fluids and plasmas. Computational fluid dynamics (CFD) simulations, which solve partial differential equations governing fluid or plasma motion on discretized grids, have been central to this progress. Recent advances have pushed the resolution and runtime of legacy numerical models to unprecedented levels while enabling newer codes to incorporate increasingly sophisticated physics. However, further scaling of these simulations has become a significant challenge, largely due to the comparatively modest improvements in networking capabilities relative to the rapid growth of floating-point performance in modern GPU-accelerated clusters. In this article, we introduce a novel load-balancing routine CRUX designed to scale efficiently for the most demanding CFD grids in astrophysics. Unlike traditional approaches based on space-filling curves, our method dynamically accounts for computational cost disparities among mesh blocks evolved with different timesteps while minimizing communication overhead. It is also able to take into account heterogeneous hardware. Through an extensive suite of benchmarks featuring up to 5,400 GPUs on OLCF Frontier and ALCF Aurora, we demonstrate that our load-balancing algorithm outperforms space-filling curve methods across all key metrics, including load uniformity, memory consumption, and communication efficiency, making it a robust solution for next-generation CFD simulations.