CAPCG 复盘（三）

应用背景

本次比赛的应用为 POP (Parallel Ocean Program)。

POP 是由 LANL 在能源部的 CHAMMP 计划赞助下开发的，该计划将大规模并行计算机引入了气候建模领域。上世纪60年代末，美国国家海洋和大气管理局地球物理流体动力学实验室的 Kirk Bryan 和 Michael Cox 首次开发了 Bryan-Cox-semtner 海洋模型，而 POP 是在这个模型的基础上进行开发的。POP 的第一个版本是由 Semtner 和 Chervin 开发的。

1. 这是一个为了研究海洋气候系统提出的三维洋流模型，适用于海洋学和气候学研究；
2. 研究人员通过实验发现，POP 模拟得到的结果总体上符合真实数据的大致分布趋势，但是仍然需要在更高分辨率下进行预测才能进一步减小与真实数据之间的误差；
3. 由于 POP 里面的 stencil 计算以及 global reduction 存在一定的访存及通信瓶颈，所以实际工作性能不能达到平台的最佳性能（计算的瓶颈较小，访存的瓶颈较大）

应用运行的大体过程如下：

• 对地图进行网格划分，每个进程分配部分任务；
• 每个进程独自计算自己的任务；
• 每个网格需要和相邻的网格交换边界信息；
• 每次迭代过程需要若干次 global reduction 收集全局信息。

在这里多说一句，海洋模型为什么会跟共轭梯度法有关联。这是因为海洋模型本质上是通过连续的微分方程描述的，而求解微分方程的常用做法是离散化，之后通过一定的变换问题就能转化为求解 $Ax=b$ 这种线性方程，也就可以使用共轭梯度法进行求解了。

应用中主要求解的是正压 (barotropic) 模型，stencil operator 的作用是差分，可以发现源码中矩阵乘法的部分都是使用 stencil 计算。这也和微分方程的离散化有所对应（可以看成稀疏矩阵乘法的一种替代形式）。

运行环境

CPU: Intel(R) Xeon(R) Gold 6140 CPU @ 2.30GHz

L3 Cache: 25344 KB

频率：2.30GHz

核心数: 18

内存: 192GB DDR4

性能分析

使用 Intel Vtune 对原始代码进行性能分析，结果如下：

可见后端载入的指令有很大部分由于访存瓶颈无法执行。

尝试直接对 PCG 进行优化

采用的算法伪代码如下：

划分的网格结果如下：

我们进行了以下优化：

• 将求解的主体部分由 Fortran 改写为 C++；
• 将原先的 global_sum 进行改写，原先是每个进程将自己要进行求和的数据放在一个数组中，调用 global_sum 之后，每个进程先对自己的数组求和，再对所有的进程求和。这样做损伤 cache 的局部性，我们将其改为每个进行在计算过程中就对自己负责的部分进行求和，存放在 local_sum 中，之后调用 MPI_Allreduce 对这个 local_sum 进行求和。改写之后速度有明显提升；
• 尽可能减少中间变量，因为访存存在瓶颈，每多一次复制拷贝的开销都会对访存造成伤害，减少中间变量之后效果提升明显。

最终 21.12 / 32.6 s.

存在的问题：

• 算法本身每次迭代过程中存在两次 global_sum，分别是 $r_k^Tz_k$ $p_k^TAp_k$ 的计算，这属于算法本身的缺陷，我们希望全局通信越少越好；
• 残差是由自身迭代计算，可能由于浮点数计算截断等出现误差，不过没有很大问题；
• 尝试了修改划分网格的方法（程序本身提供的），但看上去负载比较均衡，修改之后也没有明显提高。

尝试使用 PCSI 算法

根据参考文献，global_sum 在处理器增长之后将成为主要瓶颈，虽然我们的处理器没有这么多，但可以看出 global_sum 还是瓶颈之一。

PCSI 算法的主体迭代过程中不需要任何 global_sum，这是因为 PCG 在计算时需要很多内积计算，而内积计算就需要用到 global_sum。PCSI 是一种 Chebyshev Iteration 计算方法，这种迭代方法在求解 $Ax=b$ 是避免了内积计算，也就不需要 global_sum，但代价是需要对系数矩阵 $A$ 有一些先验知识，需要知道 $A$ 特征值的上界和下界。求解特征值范围使用的是 Lanczos 迭代法，每次迭代需要 2 次 global_sum，但因为可以设置其迭代精度，迭代次数不多，所以理论上会有所加快。

我们所做的其实很 PCG 很类似：

• 求解主体改为 C++；
• 修改 global_sum 为边计算边求和；
• 尽量减少中间变量。

最终 22.36 / 32.6 s.

遇到的问题：

• 没有想象中超过 PCG 的表现，可能是问题的规模太小，无法完全发挥算法的优势；
• 使用的系数矩阵的特征值是用 Lanczos 法迭代出来的，我们发现奇怪的是如果得到的特征值上下界越精确，PCSI 主体部分反而收敛的越慢。经过使用 MKL 等求解验证发现确实如此，不是算法实现有误。这也是我们最终放弃 PCSI 的一个重要原因，因为现场会有测试数据，如果算法不够稳定有现场翻车的可能性。

但最后的冠军使用的是 PCSI 算法，刚知道这个消息的时候我们也比较惊讶，当然有可能的原因会在后面解释。

CAPCG

算法伪代码如下：

CAPCG 的数学推导还是较为复杂，这里只阐述一下算法大概思路：

• 算法主要分为两层循环：
  • 外层循环是主迭代，负责检查是否达到最大迭代次数或者迭代精度以退出循环；
  • 由于之前的 PCG 算法是每次只搜索一个方向，可以看成只走一步，而 CAPCG 选择往前走 s 步，每次外层迭代搜索 s 个方向。
• 每次内层迭代并不直接更新 $r, x$ 等，而是更新它们在对应的搜索空间的系数，在内层迭代结束之后，系数与搜索空间的基相乘完成更新；
• 关于计算 $V_k$ 也就是代码中的 btrop_matpower_operator 是用来构建搜索空间 krylov subspace 的一组基：$V^{k}=\left[P^{k}, R^{k}\right]=\left[\rho_{0}(A) p^{s k}, \rho_{1}(A) p^{s k}, \ldots, \rho_{s}(A) p^{s k}, \rho_{0}(A) r^{s k}, \rho_{1}(A) r^{s k}, \ldots, \rho_{s-1}(A) r^{s k}\right]$，$\rho_i$ 是次数为 i 的多项式，由 $\rho_{i}(z)=\gamma_{i}\left(A-\theta_{i} I\right) \rho_{i-1}(z)+\sigma_{i} \rho_{i-2}(z)$ 三项递推式产生，凡是类似用这样的三项递推产生的多项式之后作用在一个向量上都可以看成是一个 matrix power kernel;
• 之前的内积全部用矩阵相乘代替，避免了全局通信。

改用 CAPCG 之后，刚开始就有不错的效果，迭代次数减少，速度加快，之后我们进行了如下优化：

• 使用 C++ 改写算法主体，数组访问等使用模板重写 (19.71 / 32.6 s)；
• 下调迭代精度，一些参数数组（数字大小不会改变，这样做可以提高 Cache 利用率）使用 float 存储 (13.39 / 32.6 s)；
• 去掉冗余的数组拷贝和不必要的初始化 (11.68 / 32.6 s)；
• 对算法中的一些参数进行调整，例如搜索步数等，使用脚本得到最优结果(10.32 / 32.6 s)；
• 加入一些编译选项，如：循环展开，向量化等 (9.98 / 32.6 s)。

我们的优化结果是采用这个 CAPCG 算法中最快的，但是后期优化过程中我们也发现个这个算法存在的缺陷：

• 存储 $V_k$ 的数组，每个进程有 30+M 大小，一共 24 个进程，会造成大量 L3 Cache Miss，访存有很大瓶颈；
• 由于对 $x$ 的计算是算好了直接加上去，而且对于残差的计算也是迭代进行的，是无反馈的，所以如果降低中间过程变量的计算精度（改为 float）则很容易出现不收敛的情况。

相比之下，PCSI 算法的残差每次都是通过 $x$ 计算得到，具有一定的反馈效果，因此第一名使用 PCSI 算法将精度调低（据悉有些比较激进的尝试），虽然增加了迭代次数，但每次迭代的时间却减少很多。并且 PCSI 中没有像 $V_k$ 一样对访存如此不友好的数组，其实还是有一定的优化空间。考虑到这次只是一个比赛，其实后来觉得我们没有必要完全放弃 PCSI，最后 CAPCG 的优化进入瓶颈期之后，回过头来将 CAPCG 的一些优化方法和想法应用到 PCSI 上，也许可能有当时意想不到的收获。

参考文献

1. A Performance Model of the Parallel Ocean Program.

2.Y. Hu et al., “Improving the scalability of the ocean barotropic solver in the community earth system model,” SC ‘15: Proceedings of the International Conference for High Performance Computing, Networking, Storage and Analysis, Austin, TX, 2015, pp. 1-12.

3.M. Hoemmen, Communication-avoiding Krylov subspace methods, PhD thesis, University of California, Berkeley, 2010.

4.Erin Carson and James Demmel, A Residual Replacement Strategy for Improving the Maximum Attainable Accuracy of s-step Krylov Subspace Methods. SIAM J. Matrix Anal. Appl., 35(1), pp. 22-43.