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

TL;DR

This paper enhances a lattice QCD solver to efficiently handle multiple right-hand sides across different architectures, optimizing data layouts and SIMD utilization, and analyzing performance portability and architecture-specific benefits.

Contribution

It introduces a flexible data layout and SIMD optimization for multi-RHS lattice QCD solvers, along with a comprehensive performance analysis across architectures.

Findings

Achieved similar speedups on x86 and Arm clusters.

Demonstrated performance portability of the optimizations.

Provided insights into architecture-specific performance factors.

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-$\alpha$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.