Performance-Portable Optimization and Analysis of Multiple Right-Hand Sides in a Lattice QCD Solver

TL;DR

The paper tackles memory-bound challenges in lattice QCD by extending the DD-$\alpha$AMG solver to handle multiple rhs through rhs blocking in both the Wilson-Dirac operator and a batched GMRES solver. It introduces SIMD-friendly data layouts and investigates Arm SME implementations to widen portability across HPC platforms. Across x86 (JUWELS) and Arm (Ookami, HAICGU), the study reports 10–24% kernel speedups with rhs blocking, while highlighting that real gains are constrained by memory bandwidth, cache behavior, and compiler optimizations; an emulator assesses SME benefits under idealized conditions. The work demonstrates performance portability and provides a nuanced analysis of architecture-specific factors that govern speedups, offering guidance for future optimizations and SME-oriented developments in LQCD. The findings have practical impact for deploying memory-bound solvers on heterogeneous HPC systems and for leveraging emerging Arm technologies in scientific computing.

Abstract

Managing the high computational cost of iterative solvers for sparse linear systems is a known challenge in scientific computing. Moreover, scientific applications often face memory bandwidth constraints, making it critical to optimize data locality and enhance the efficiency of data transport. We extend the lattice QCD solver DD-$α$AMG to incorporate multiple right-hand sides (rhs) for both the Wilson-Dirac operator evaluation and the GMRES solver, with and without odd-even preconditioning. To optimize auto-vectorization, we introduce a flexible interface that supports various data layouts and implement a new data layout for better SIMD utilization. We evaluate our optimizations on both x86 and Arm clusters, demonstrating performance portability with similar speedups. A key contribution of this work is the performance analysis of our optimizations, which reveals the complexity introduced by architectural constraints and compiler behavior. Additionally, we explore different implementations leveraging a new matrix instruction set for Arm called SME and provide an early assessment of its potential benefits.

Figures (6)

• Figure 1: Example of matrix operation using two data layouts with $b=4$, where dotted lines and arrows represent the storage order in memory.
• Figure 2: Demonstration of $A_{1^{st}col} \otimes M_{1^{st}row}$ using SME.
• Figure 3: Performance per node scaling with arithmetic intensity in the Wilson-Dirac operator evaluation for the $128\times64^3$ lattice on JUWELS and Ookami, where the percentage labels denote architectural efficiencies.
• Figure 4: Single-node performance comparison of the Wilson-Dirac operator evaluation for the $64\times 16^3$ lattice across JUWELS, Ookami, and HAICGU.
• Figure 5: Performance per node scaling with arithmetic intensity for 100 iterations of batched GMRES for the $128\times64^3$ lattice on JUWELS.