1 引言

强相互作用是自然界中4种(强、弱、电磁、引力)基本相互作用之一. 1974年, Wilson建立的格点量子色动力学(Lattice quantum chromodynamics, 格点QCD)规范理论^[1]是从第一性原理出发研究强相互作用的基本理论和方法, 尤其是在高能物理的低能区, 格点QCD是当前已知的最为系统最为可靠的非微扰理论方法, 对高能物理实验和理论的研究都有着重要意义.

格点QCD算法与软件的研究是利用该方法开展强子谱学、核子结构、QCD相变、标准模型精确检验和新物理寻等高能物理研究的基础支撑.

在国际上, 格点QCD作为重要的高性能计算应用, 并且形成诸多优秀的模拟计算软件. 例如21世纪初美国USQCD的Chroma^[2], 以及2010年发展出的GPU异构加速计算库QUDA^{[3, 4]}; 日本从2009年开始研发适应日本格点计算研究需求的计算软件Bridge++^{[5, 6]}, 而后发展出面向日本超级计算机“富岳”的格点QCD的计算库QWS^{[7, 8]}; 近年来英国的Grid^{[9, 10]}在数据处理的矢量化设计上表现突出, 对于不同架构的机器具有较好适应性.

目前, 随着国内高性能计算机的发展, 国内格点QCD研究团队积极开展面向国产超算平台的格点计算软件研发工作. 在申威、天河等不同超算平台上探索式地进行了软件移植和自主开发工作^[11-16].

本团队致力于使用HIP框架开发针对国产ROCm异构平台的格点QCD异构计算程序库Openchiral, 首次在HIP框架下自主实现了国产ROCm环境的格点QCD计算所需基本核心功能模块. 格点QCD计算的核心热点模块为有限差分的狄拉克(Dirac)方程求解, 也即QCD传播子的求解过程, 下称求解器, 其热点运算为Dslash运算(参见第2.1节). 本文工作首先实现了格点QCD相关的矩阵向量操作、基本核函数Dslash和共轭梯度(conjugate gradient, CG)求解器出发, 而后逐步进行算法改进和计算优化, 实现了在ROCm国产异构平台上高效稳定的格点QCD传播子求解器.

2 格点QCD求解器及基本流程

格点QCD是定义在 ${N_x} \times {N_y} \times {N_{\textit{z}}} \times {N_t}$ 四维结构网格上的理论和方法. 在格点上定义费米子场 $ \varphi (\tilde x) $ , 格点之间链接上定义规范场 $ {U_\mu }(\tilde x) $ , 这里 $ \tilde x $ 表示四维时空坐标, $ \mu $ 表示四维时空方向. 每个网格点是一个张量点, 是由QCD色空间和旋量空间张成的3 $ \otimes $ 4维空间. 计算的核心热点模块为有限差分的狄拉克方程求解, 也即QCD传播子的求解过程, 下称求解器. 其形式上类似于有限差分方程 $Ax = b$ 的求解.

2.1 Dslash运算

本文以基本的格点QCD的Wilson费米子作用量为例, 对应的求解方程形式如下:

$ \left[ {{\gamma _\mu }{D_\mu } + m} \right]\varphi = M\varphi = {\text{b}} $

(1)

其中, $M$ 称为费米子矩阵, 是一个规模庞大的稀疏矩阵, 其维数规模为 ${N_x} \times {N_y} \times {N_{\textit{z}}} \times {N_t} \times 12$ ; 费米子矩阵 $M$ 具体形式为:

$ \begin{split} M = & \left( {4 + m} \right){\delta _{\tilde x, \tilde y}} - \frac{1}{2}\sum\limits_{\mu = 0}^3 {\left[ {(1 - {\gamma _\mu }){U_\mu }(\tilde x){\delta _{\tilde x + \mu , \tilde y}}} \right.} \\ & + (1 + {\gamma _\mu })U_\mu ^{\dagger} (\tilde x - \mu )\left. {{\delta _{\tilde x - \mu , \tilde y}}} \right] \end{split} $

(2)

其中, $ {\gamma _\mu } $ 是确定的四维狄拉克矩阵, $ {U_\mu }(\tilde x) $ 是 $SU(3)$ 矩阵. 由式(2)可知, 该计算是四维时空结构网格上的四维9点的模板(stencil)计算, 点与点之间的运算涉及 $3 \otimes 4$ 维张量计算. 根据如式(2)所示的稀疏矩阵 $M$ 的数学形式, 为了增加数据复用、减少内存占用、降低访存压力, 实现中不用直接存储高维的 $M$ 矩阵, 而是在计算过程中利用组态数据 ${U_\mu }(x)$ 构建费米子矩阵乘向量操作 $M\varphi $ , 称为Dslash操作, 该运算是迭代求解中核心热点运算.

2.2 格点QCD求解器

稳定高效的Dslash运算和求解器将会显著提升格点QCD计算的整体性能. 通常使用的共轭梯度等迭代求解算法需调用第2.1节所述的Dslash操作, 现有的稀疏矩阵求解库难以满足需求, 通常需要为其专门实现一套求解器.

格点QCD求解器模块工作流程如下:

1) 初始化费米子数据.

2) 加载组态数据.

3) 开始迭代求解.

4) 对求出的解进行验算.

计算流程图如图1所示, 待求解矩阵 $M$ 由组态数据与费米子场数据构建, 不直接存储.

程序中主要包含以下模块:

1) 数据初始化模块: 初始化费米子场、组态数据.

2) 迭代求解模块: 使用费米子场、组态数据构造所求的矩阵方程 $Mx = b$ , 并对其进行迭代求解.

3) Dslash模块: 在异构加速卡完成迭代求解过程中的矩阵向量乘运算.

4) 验算模块: 对求解器给出的解进行验算, 确认正确性.

3 优化方案设计与实现

在基于HIP框架实现面向国产ROCm异构平实现核心运算Dslash基础上, 进一步通过算法改进和异构访存优化, 用以提高收敛速度、降低问题规模. 具体优化方案如下:

1)在实现CG (算法1)求解器的基础之上实现BiCGSTAB (算法2)求解器, 相对于CG求解器迭代次数更少、收敛速度更快.

2)采用奇偶预处理方法降低求解问题的规模, 进一步减少迭代次数并提高每轮迭代的速度.

3)在异构访存优化方面, 通过Dslash合并异构访存与改用锁页内存方式, 减少每轮迭代时CPU与加速卡的通讯开销, 降低每轮迭代耗时.

算法1. CG算法

procedure CG ( $\scriptstyle{{M^\dagger }M}$ , x, b) $\scriptstyle{{r_0}: = {M^\dagger }b - {M^\dagger }M{x_0}}$ $\scriptstyle{{p_0}: = {r_0}}$ for i=0, 1, 2, …, do $\scriptstyle{{\alpha _i}: = \frac{{{r_i} \cdot {r_i}}}{{\left( {{M^\dagger }M{p_i}} \right) \cdot {p_i}}}}$ /*计算步长*/ $\scriptstyle{{x_{i + 1}}: = {x_i} + {\alpha _i}{p_i}}$ /*修正解向量*/ $\scriptstyle{{r_{i + 1}}: = {r_i} - {\alpha _i}{M^\dagger }M{p_i}}$ /*计算残差向量*/ if $\scriptstyle{\left| {\left| {{r}} \right|} \right| < {\varepsilon }}$ then break /*判断是否收敛*/

end if $\scriptstyle{{\beta _i}: = \frac{{{r_{i + 1}} \cdot {r_{i + 1}}}}{{{r_i} \cdot {r_i}}}}$ $\scriptstyle{{p_{i + 1}}: = {r_{i + 1}} + {\beta _i}{p_i}}$ /*计算梯度方向*/ end for $\scriptstyle{x: = {x_{i + 1}}}$ return $\scriptstyle{x}$ end procedure

算法2. BiCGSTAB算法

procedure BiCGSTAB (M,x, b) $\scriptstyle{{r_0}: = b - M{x_0}}$ $\scriptstyle{r_0^{'}: = {r_0}}$ $\scriptstyle{{p_0}: = {r_0}}$ for i=0, 1, 2, …, do $\scriptstyle{{a_i}: = \frac{{{r_i} \cdot r_0^{{'}}}}{{\left( {M{p_i}} \right) \cdot r_0^{{'}}}}}$ /*计算步长*/ $\scriptstyle{{s_i}: = {r_i} - {\alpha _i}M{p_i}}$ $\scriptstyle{{\omega _i}: = \frac{{\left( {M{s_i}} \right) \cdot {s_i}}}{{\left( {M{s_i}} \right) \cdot \left( {M{s_i}} \right)}}}$ $\scriptstyle{{x_{i + 1}}: = {x_i} + {\alpha _i}{p_i} + {\omega _i}{s_i}}$ /*修正解向量*/ $\scriptstyle{{r_{i + 1}}: = {s_i} - {\omega _i}M{s_i}}$ /*计算残差向量*/ if $\scriptstyle{\left| {\left| {{r}} \right|} \right| < {\varepsilon }}$ then break /*判断是否收敛*/ end if $\scriptstyle{{\beta _i}: = \frac{{{\alpha _i}}}{{{\omega _i}}} \times \frac{{{r_{i + 1}} \cdot r_0^{{'}}}}{{{r_i} \cdot {r_0}}}}$ $\scriptstyle{{p_{i + 1}}: = {r_{i + 1}} + {\beta _i}\left( {{p_i} - {\omega _i}M{p_i}} \right)}$ /*计算梯度方向*/ end for $\scriptstyle{x: = {x_{i + 1}}}$ return $\scriptstyle{x}$ /*返回结果*/end procedure

3.1 提高收敛速度: 实现BiCGSTAB求解器

相对于CG算法(算法1), BiCGSTAB算法(算法2)能够求解非正定矩阵. 更重要的是, BiCGSTAB可以减少迭代次数, 提升计算效率. 在第i轮迭代中, 记 ${\rho _i} = \left( {I - {\omega _i}M} \right)$ , $r_i^B$ 、 $r_i^G$ 分别为BiCGSTAB、CG算法的残差, 则有:

$ \begin{split} r_i^B &= {\rho _i}\left( M \right)r_i^G = \left( {I - {\omega _i}M} \right){\rho _{i - 1}}\left( M \right)r_i^G \\ &= \left( {I - {\omega _i}M} \right)\left( {I - {\omega _{i - 1}}M} \right){\rho _{i - 2}}\left( M \right)r_i^G \\ &= \cdots = \mathop \prod \limits_{k = 0}^i \left( {I - {\omega _{\text{k}}}M} \right)r_i^G \end{split} $

(3)

由式(3)可知当 $r_i^G = 0$ 时, $r_i^B = 0$ , 表明二者具有相同的收敛性与精度, 额外的相乘操作 $\;{\rho _i}\left( M \right)$ 使得BiCGSTAB每轮迭代更加平滑, 算法收敛到0的速度更快, 迭代次数下降.

在迭代求解过程中, 主要的热点是Dslash操作. BiCGSTAB需要计算 $M{p_i}$ 、 $M{s_i}$ , 故每轮迭代共需进行两次Dslash操作. CG需计算 $M{p_i}$ , 但由于格点QCD的费米子矩阵 $M$ 不保证正定, 使用CG求解时方程两侧需要左乘 ${M^\dagger }$ . 以保证待求矩阵的正定性, 即 ${M^\dagger }Mx = {M^\dagger }b$ , 所以CG和BiCGSTAB每轮迭代均需要进行两次Dslash操作. 由此可以得知, BiCGSTAB每轮迭代开销与CG大致相同, 求解速度的提升主要来自更少的迭代次数.

3.2 降低问题规模: 奇偶预处理

根据格点QCD离散化的时空结构网格和费米子矩阵类差分形式, 对费米子矩阵 $M$ 和对应的费米子场向量做奇偶预处理(even-odd precondition), $M$ 矩阵会可以划分为4块子矩阵, 即:

$ M = \left[ {\begin{array}{*{20}{c}} {{M_{ee}}}&{{M_{eo}}} \\ {{M_{oe}}}&{{M_{oo}}} \end{array}} \right] $

(4)

其中, ${M_{ee}}$ 和 ${M_{oo}}$ 是实对角部分, ${M_{eo}}$ 和 ${M_{oe}}$ 是非对角部分. 奇偶预处理前后矩阵元分布如图2所示.

经过奇偶预处理后的费米子矩阵M, 可以做如下形式的分解:

$ M = L\tilde MU $

(5)

其中, $ \tilde M $ 为块对角矩阵.

$ \tilde M = \left[ {\begin{array}{*{20}{c}} {{M_{ee}}}&0 \\ 0&{{M_{oo}} - {M_{oe}}M_{ee}^{ - 1}{M_{eo}}} \end{array}} \right] $

(6)

因此, 经过奇偶预处理后, 式(1)所示的线性求解问题 $M\varphi = b$ 变为如下形式:

$\left\{ { \begin{gathered} {L}{\tilde MU}{\varphi} ={ b} \\ {\tilde M}{\varphi }' = {b'} \end{gathered}} \right. $

(7)

其中, $b' = {L^{ - 1}}b$ , $\varphi = {U^{ - 1}}\varphi '$ . L和U形式分别如下:

$ L = \left[ {\begin{array}{*{20}{c}} {1}&{0} \\ {{M_{oe}}M_{ee}^{ - 1}}&{1} \end{array}} \right], \; U = \left[ {\begin{array}{*{20}{c}} {1}&{M_{ee}^{ - 1}{M_{eo}}} \\ {0}&{1} \end{array}} \right] $

其逆矩阵 ${L^{ - 1}}$ 和 ${U^{ - 1}}$ 很容易从理论上推导得到而不需要对L和U进行数值求逆, 其形式为:

$ {L^{ - 1}} = \left[ {\begin{array}{*{20}{c}} I&0 \\ { - {M_{oe}}M_{ee}^{ - 1}}&I \end{array}} \right], \; {U^{ - 1}} = \left[ {\begin{array}{*{20}{c}} I&{ - M_{ee}^{ - 1}{M_{eo}}} \\ 0&I \end{array}} \right] $

其中, ${M_{ee}}$ 和 ${M_{oo}}$ 均为对角矩阵.

$ M_{oo}^{ - 1} = M_{ee}^{ - 1} = 1/\left( {4 + m} \right) $

因此, 经过奇偶预处理后, 需要数值迭代求解仅为式(6)和式(7)下半部分, 仍记为 ${\tilde M}{\varphi }' = {b'}$ . $\tilde M$ 的维度是 $M$ 的1/2, 因此降低了问题规模, 可以有效减少迭代次数. 同时, 由于 $\tilde M$ 减小, 每次Dslash操作的开销也大幅降低, 能够降低每轮迭代的耗时.

3.3 异构访存优化 3.3.1 Dslash操作合并异构访存

在迭代求解过程中, Dslash操作是计算热点, 提升这部分的计算效率可以显著提升迭代求解的速度. 我们基于HIP框架实现了Dslash的异构计算加速, 优化了异构计算访存.

经过奇偶预处理后, 迭代求解中的Dslash操作的内容就变为了 ${M_{oe}}{M_{eo}}x$ , 也就是连续两次Dslash操作, ${M_{oe}}$ 、 ${M_{eo}}$ 分别是同一矩阵 $M$ 的非对角区域. 每次Dslash操作会在加速卡显存、主存之间传输两次, 如图3(a)所示. 令 $v = {M_{eo}}x$ , $w = {M_{oe}}v$ , 则第1次Dslash操作 $v = {M_{eo}}x$ 需要从主存传输组态U、向量 $x$ ; 从显存传回计算结果向量 $v$ . 第2次Dslash操作需要从主存传输组态U、向量 $v$ ; 从显存传回计算结果向量 $w$ . 故连续两次Dslash操作就要在主存、显存间传输4次.

由于费米子矩阵元素数量是所依赖数据元素数量的12倍, 所以Dslash操作不直接构建和存储稀疏矩阵 $M$ , 而是采用模板计算, 在计算时使用组态数据U生成待计算的矩阵元素, 故前后两次Dslash 所需的组态数据U相同. 又因连续Dslash操作上一次运算的结果即为下一次运算的输入, 因此可以将4次访存合并为2次: 图3 (b), 当前一次Dslash操作结束后, 将两次计算共同依赖的数据U和前一次Dslash操作 $v = {M_{eo}}x$ 的计算结果存储在显存上, 作为下一次Dslash 操作的输入数据, 节约了2次访存操作.

3.3.2 提高异构访存速度

迭代求解时, 每轮迭代进行Dslash操作时都需要通过PCIe接口向加速卡传输组态数据与费米子场数据, 这一操作是由HIP框架提供的API完成的.

默认情况下C、C++中使用new或malloc所分配的内存空间是可分页内存. 在HIP框架中, 从主存向显存传输数据前必须先分配临时页锁定数据, 再传输到显存, 如图4(a)所示.

HIP框架提供了分配锁页内存的功能, 即hipHostMalloc()函数. 相比可分页内存, 将待传输的数据存入锁页内存省去了锁页与拷贝的步骤, 可以加快数据传输速度, 如图4(b)所示. 在格点QCD计算迭代求解时, 需要传输的数据为费米子场向量与规范场数据. 将费米子与规范场的数据存储在锁页内存, 可以大幅提高加速卡与内存通讯的速度.

4 实验分析

本工作完成了以上几种优化手段的代码实现与性能测试. 在中科先导一号计算平台上, 优化后的求解程序相比最初的程序大体达到了30倍的加速比.

为适应不同问题规模, 本论文以下性能测试为弱扩展性测试, 具体测试规模如表1所示: 每节点的负载不变, 总计算量随节点数成比例扩展. 程序测试的规模以格子的 $x$ 、 $y$ 、 $\textit{z}$ 、 $t$ 指标体现. 每张加速卡承担的计算任务规模为子格子大小, 设为 ${16^4}$ ; 总格子规模为总计算任务规模, 随加速卡数增加. 令加速卡数为 ${2^c}$ , 总格子规模为 $16 \times {2^x}$ 、 $16 \times {2^y}$ 、 $16 \times {2^{\textit{z}}}$ 、 $16 \times {2^t}$ , 则有 $c = x + y + {\textit{z}} + t$ .

以 $32$ 卡为例: $32 = {2^5}$ , 子格子为 ${16^4}$ , $5 = 0 + 1 + 2 + 2$ , 则总格子各维度规模为 $16 \times {2^0}$ 、 $16 \times {2^1}$ 、 $16 \times {2^2}$ 、 $16 \times {2^2}$ , 即 $16 \times 32 \times 64 \times 64$ .

4.1 求解速度优化效果

在表2中, 我们给出了不同规模下上述几种优化手段的求解耗时表现, 表格中从左到右优化手段顺次叠加, “锁页内存”即为最终优化后的程序. 可以看出, 充分优化后的程序在不同规模都达到了30倍左右的加速比. 图5直观地展示、对比了几种优化手段之间的耗时表现.

在表3中, 给出各优化手段与对应减少的耗时, 从各个方法所优化耗时与最终减少的耗时的占比来看, 奇偶预处理降低了所求解问题的规模, 对耗时优化贡献最大, 平均贡献总优化耗时的48%; BiCGSTAB算法减少了迭代次数, 平均贡献了总优化耗时的26%; 合并访存、锁页内存降低了加速卡访存的开销, 降低了每轮迭代的耗时, 对总优化耗时分别贡献了12%与13%.

4.2 迭代次数优化效果

本文所采取的几种优化手段中, BiCGSTAB与奇偶预处理有效减少了收敛迭代次数, 接下来将两者与原始程序的迭代收敛表现做对比分析.

在表4中, 我们在不同规模下对比了几种收敛优化手段的迭代收敛表现, 从左到右优化手段顺次叠加, 分别为原始程序(CG)、BiCGSTAB求解器与奇偶预处理的BiCGSTAB求解器. 从中我们可以看出, 原始程序中的CG求解器, BiCGSTAB与奇偶预处理在不同规模下都展现出了大致相同的优化效果. 最终优化后的程序求解收敛所需的迭代次数在不同规模下平均减少为原始程序的33.4%.

在图6中, 我们给出了3种算法在128卡下求解 $32 \times 32 \times 64 \times 128$ 时的收敛表现, 纵轴为求得残差, 当残差为0时即求解完毕. 曲线越陡峭、越靠左说明算法收敛得越快. 从图中可以直观地看出BiCGSTAB与奇偶预处理在迭代过程中的收敛情况.

5 结论与展望

本文针对格点QCD求解器设计实现, 采取了多项优化措施, 实现了BiCGSTAB求解器, 显著减少迭代次数; 采用了奇偶预处理技术, 降低了求解问题的规模; 针对HIP/ROCM异构计算平台采取了合并、加速访存等优化方式. 本文还开展了多种规模的并行测试, 最大规模为64节点、256块加速卡, 相比原始程序实现了约30倍的加速比.

格点QCD是一种重要的高性能计算应用, 性能优化后的程序可以更好地支撑格点QCD研究, 充分发挥国产异构超算平台的算力.