Enhancing data locality of the conjugate gradient method for high-order matrix-free finite-element implementations

TL;DR

This paper introduces a data locality enhancement for the conjugate gradient method in high-order finite-element computations, significantly reducing memory transfers and doubling performance on large-scale parallel systems.

Contribution

It presents a novel variant of the CG method that improves data locality through vector update interleaving and enumeration, leading to substantial performance gains.

Findings

Achieves around 90% vector entry transfers from slow RAM per iteration.

Doubles the speed compared to standard and pipelined CG methods on large-scale systems.

Maintains similar performance near the strong scaling limit.

Abstract

This work investigates a variant of the conjugate gradient (CG) method and embeds it into the context of high-order finite-element schemes with fast matrix-free operator evaluation and cheap preconditioners like the matrix diagonal. Relying on a data-dependency analysis and appropriate enumeration of degrees of freedom, we interleave the vector updates and inner products in a CG iteration with the matrix-vector product with only minor organizational overhead. As a result, around 90% of the vector entries of the three active vectors of the CG method are transferred from slow RAM memory exactly once per iteration, with all additional access hitting fast cache memory. Node-level performance analyses and scaling studies on up to 147k cores show that the CG method with the proposed performance optimizations is around two times faster than a standard CG solver as well as optimized pipelined CG and s-step CG methods for large sizes that exceed processor caches, and provides similar performance near the strong scaling limit.