In-Situ Assessment of Device-Side Compute Work for Dynamic Load Balancing in a GPU-Accelerated PIC Code

TL;DR

This paper introduces GPU-specific strategies for assessing compute work to improve dynamic load balancing in GPU-accelerated PIC codes, significantly enhancing performance on large-scale GPU systems.

Contribution

It presents novel GPU-amenable methods for in-situ compute work assessment and demonstrates their effectiveness in optimizing load balancing for large-scale GPU applications.

Findings

Achieved 62-74% of theoretical maximum speedup on Summit with 6144 GPUs.

Improved load balancing yields 3.8x speedup over static methods on 96 GPUs.

Optimal data collection strategies for GPU compute work assessment identified.

Abstract

Maintaining computational load balance is important to the performant behavior of codes which operate under a distributed computing model. This is especially true for GPU architectures, which can suffer from memory oversubscription if improperly load balanced. We present enhancements to traditional load balancing approaches and explicitly target GPU architectures, exploring the resulting performance. A key component of our enhancements is the introduction of several GPU-amenable strategies for assessing compute work. These strategies are implemented and benchmarked to find the most optimal data collection methodology for in-situ assessment of GPU compute work. For the fully kinetic particle-in-cell code WarpX, which supports MPI+CUDA parallelism, we investigate the performance of the improved dynamic load balancing via a strong scaling-based performance model and show that, for a laser-ion acceleration test problem run with up to 6144 GPUs on Summit, the enhanced dynamic load balancing achieves from 62%--74% (88% when running on 6 GPUs) of the theoretically predicted maximum speedup; for the 96-GPU case, we find that dynamic load balancing improves performance relative to baselines without load balancing (3.8x speedup) and with static load balancing (1.2x speedup). Our results provide important insights into dynamic load balancing and performance assessment, and are particularly relevant in the context of distributed memory applications ran on GPUs.