面向 GPU 集群各向异性 PDE 的 PB 级椭圆求解器

内存受限型应用（如大型稀疏方程组求解器）对 GPU 而言仍是一项挑战。快速求解器应基于数值高效的算法，并以最小化全局内存访问的方式实现。为求解高达一万亿（$10^{12}$）未知数的系统，代码必须在大型 GPU 集群上高效利用数百万个独立的处理器核心。我们描述了两种算法最优的迭代求解器在多 GPU 上的实现，这些求解器用于大气建模中遇到的各向异性椭圆偏微分方程（PDE）。在此应用中，条件数很大但与网格分辨率无关，且两种方法均为渐近最优，尽管其绝对性能有所不同。我们对求解器进行了并行化处理，并使其适应 GPU 架构的特定特征，特别关注高效的全局内存访问。在 16384 个 GPU 上求解具有 $0.55\cdot 10^{12}$ 未知数的方程时，我们实现了高达 0.78 PFLOPs 的性能；这相当于该机器理论峰值性能的约 $3\%$，并且在使用共轭梯度（CG）求解器时利用了超过 $40\%$ 的峰值内存带宽。尽管另一种求解器（几何多重网格算法）在每秒浮点运算次数（FLOPs）方面的性能略差，但总体而言它更快，因为它收敛所需的迭代次数更少；该多重网格算法可以在约一秒内求解具有五千亿未知数的线性 PDE。