Accelerating the Lagrangian simulation of water ages on distributed, multi-GPU platforms: The importance of dynamic load balancing

Water age is a fundamental descriptor of the source, storage, and mixing of water in watersheds. The Lagrangian, particle tracking, approach is a powerful tool for physically-based modeling of water age distributions, but its application has been hampered because it is computationally demanding. Here, we present a parallel approach for particle tracking simulations. This approach uses a multi-GPU with MPI parallelism based on domain decomposition. An inherent challenge of the distributed parallelization of Lagrangian approaches is the disparity in computational work or load imbalance (LIB) among different processing elements (PEs). In this study, load balancing (LB) schemes were proposed to dynamically balance the distribution of particles across PEs during runtime. In the followed hillslope simulations, LIB was observed in all LB-disabled runs, e.g., with a load ratio of 4.3 by using 2-GPU in the test case. LB schemes then accurately balanced the load distribution and improved the parallel scaling. Additionally, the parallel approach showed an excellent overall speedup: a 25-fold improvement using 4-GPU relative to 128 OpenMP threads. A regional-scale application further demonstrated the LB performance. The wall-clock time used by 8-GPU without LB was reduced by 31.33% after the LB was activated. Increasing 8-GPU with LB to 16-GPU with LB showed parallel scalability by reducing the wall-clock time by ∼50%. This work shows how massively parallel computing can be applied to particle tracking in water age simulations. It also demonstrates the practical importance of load balancing in this context, which enables large-scale simulations with the increased complexity of flow paths.