HIP-DFPT: Scalable Optimization of Irregular Workloads in Quantum Perturbation on GPU Clusters
Irregular workloads challenge GPUs with small tasks, unpredictable branching, and uneven load distribution. Quantum perturbation theory and Density Functional Theory (DFT) represent a prominent example, where such irregular patterns arise in large-scale simulations of materials and molecules. Despite recent efforts to accelerate these simulations on GPUs, State-of-the-Art (SOTA) methods still suffer from low GPU utilization and severe load imbalance, which hinder efficiency and scalability. In this paper, we redesign the expression of the data structure and propose an efficient implementation for quantum perturbation named HIP-DFPT. A series of general GPU optimizations are proposed, including micro-task composition strategy, GPU memory management, and multi-stream pipelines. We present a branch-dominant hybrid online performance model to predict quickly and accurately with only 5% sampled data. Furthermore, we introduce a multi-level load balancing strategy guided by performance models for both inter-GPU and intra-GPU task distribution. Our experimental results on a large-scale GPU cluster show significant improvements in computational efficiency and scalability. Using up to 8192 GPUs for simulations with up to 200,006 atoms with all-electron precision, we achieve an average load balancing efficiency of 1.06, a 55% increase in GPU occupation, and a 70% enhancement in throughput.