Efficient multi-GPU implementations of preconditioned conjugate gradient method

Tailai Ma

doi:10.1016/j.advengsoft.2025.103936

2025-06-02

PRE

DOI：10.1016/j.advengsoft.2025.103936

原文链接

原文求助

摘要

En 中文

许多数值方法依赖于求解稀疏线性系统，例如有限元分析和计算流体动力学。在许多模拟中，求解稀疏线性系统已成为瓶颈，常常消耗总模拟时间的显著部分。预条件共轭梯度（PCG）方法是求解大型稀疏线性系统的一种高效迭代方法，并且非常适合GPU并行化。随着待解决问题规模的增加，提高PCG求解器的效率至关重要。本文提出了两种高效的PCG方法实现，旨在进一步加速其在GPU上的计算性能。首先，我们使用块存储格式优化了PCG中最耗时的操作SpMV（稀疏矩阵-向量乘法），以实现高效的全局内存访问。为此，我们引入了两种新颖的SpMV算法：warp-fixed-mma和warpbalanced-mma。接下来，我们提出了两种PCG实现：PCG-cuda-graphs和PCG-persistent-kernel。这两种实现均使用了所提出的SpMV算法。在SuiteSparse Matrix Collection的2904个矩阵上的实验结果表明，我们性能最佳的SpMV算法在FP64精度下，在A100上分别比四种最新的SpMV方法cuSPARSE、DASP（Specific Dense Matrix Multiply-Accumulate Units Accelerated General Sparse Matrix-Vector Multiplication）、CSR5和TileSpMV平均快2.2倍、3.1倍、2.48倍和3.20倍。与PETSc（Portable, Extensible Toolkit for Scientific Computation）中的PCG实现相比，两种提出的PCG实现分别在NVIDIA A100 GPU和NVIDIA H100 GPU上对21个代表性矩阵实现了平均1.41倍和1.63倍、1.43倍和1.66倍的速度提升。此外，对于大型矩阵，我们提出的基于多GPU并行的PCG-persistent-kernel实现与单GPU实现相比，实现了显著的速度提升。

Keyword:

PCG

GPU

Tensor core

SpMV

Multi-GPU