金之雁，楊 磊，林雋民，王 哲

1.中國(guó)氣象科學(xué)研究院，北京100081

2.英特爾中國(guó)有限公司，北京100081

1 引言

廣義共軛余差法（Generalized Conjugate Residual method，GCR）[1]是一種用于求解非對(duì)稱線性方程組的有效算法。中國(guó)氣象局的新一代“全球區(qū)域一體化數(shù)值預(yù)報(bào)模式（Global/Regional Assimilation and PrEdiction System，GRAPES）”[2]，其動(dòng)力框架的核心是亥姆霍茲方程求解器，求解器所采用的就是GCR算法。和其他類似算法一樣，GCR算法存在的一個(gè)制約可擴(kuò)展性的問(wèn)題就是密集的全局通信。隨著CPU單節(jié)點(diǎn)運(yùn)算能力的不斷提高，這個(gè)問(wèn)題已經(jīng)成為了制約整個(gè)數(shù)值模式速率提高的主要瓶頸。

一般來(lái)說(shuō)，算法的開(kāi)銷可以歸結(jié)為兩部分：數(shù)據(jù)的計(jì)算和數(shù)據(jù)的移動(dòng)。其中數(shù)據(jù)的移動(dòng)包括節(jié)點(diǎn)內(nèi)部的移動(dòng)和節(jié)點(diǎn)間通過(guò)網(wǎng)絡(luò)的移動(dòng)。要減少這種密集的數(shù)據(jù)移動(dòng)帶來(lái)的開(kāi)銷，加州大學(xué)伯克利分校的Demmel教授最早于2007年提出了通信避免（Communication Avoiding，CA）算法[3]。CA算法的思路是：構(gòu)造能包含迭代過(guò)程中的所有長(zhǎng)向量的Krylov子空間，利用短向量的迭代替代長(zhǎng)向量的迭代，進(jìn)而避免在迭代過(guò)程中進(jìn)行通信。該思路具體應(yīng)用于對(duì)不同的算法的改造，會(huì)產(chǎn)生相應(yīng)的不同的CA算法。目前針對(duì)GCR算法，相關(guān)工作仍是空白。

數(shù)值預(yù)報(bào)模式的運(yùn)算速率是提高預(yù)報(bào)分辨率的客觀前提，“神威太湖之光”、“曙光派”等高性能計(jì)算集群的部署已經(jīng)使GCR算法中的全局通信成為提升GRAPES模式整體速率的最大瓶頸。為此，本文借鑒Demmel教授提出的思路，首創(chuàng)性地對(duì)GCR算法進(jìn)行CA改造，提出CA-GCR算法。新算法的通信次數(shù)較之原算法降低了一個(gè)數(shù)量級(jí)，同時(shí)，計(jì)算量也有一定減少。并且在中國(guó)氣象局最新部署的“曙光派”計(jì)算集群上進(jìn)行了大規(guī)模測(cè)試（最大規(guī)模16 384進(jìn)程），在可擴(kuò)展性（本文中所有“可擴(kuò)展性”均指“強(qiáng)可擴(kuò)展性”）和運(yùn)算速率上表現(xiàn)出了相對(duì)于原算法巨大的優(yōu)勢(shì)，最高達(dá)到原算法3倍的加速比。

2 相關(guān)工作

CA算法由Demmel教授提出，之后又由其本人及學(xué)生具體應(yīng)用于共軛梯度法、廣義最小剩余法等經(jīng)典算法的改造，多數(shù)表現(xiàn)出了對(duì)通信量和計(jì)算量的雙重作用，即同時(shí)減少兩者的開(kāi)銷。Demmel教授的學(xué)生Grigori于2008年將CA算法應(yīng)用于一般高斯消元法[4]；Demmel于2010年對(duì)LU分解進(jìn)行CA優(yōu)化[5]；Hoemmen于2010年對(duì)廣義最小剩余法進(jìn)行CA優(yōu)化[6]；Maryam于2013年在圖形處理單元方面實(shí)踐CA算法[7]……教授及其學(xué)生的研究成果現(xiàn)已涵蓋多數(shù)主流算法，但GCR算法的相關(guān)工作仍是空白，目前國(guó)內(nèi)外都沒(méi)有相應(yīng)方案。

3 算法描述

3.1 原始算法

2008年曹建文等人對(duì)GRAPES模式中亥姆霍茲方程的系數(shù)矩陣進(jìn)行ILU分解，形成了“帶預(yù)條件的廣義共軛余差法”（Preconditioned Generalized Conjugate Residual method，P-GCR）[8]，目前已應(yīng)用于業(yè)務(wù)運(yùn)行。在本文中，以該算法為原始算法。其具體算法如下：

算法1帶預(yù)條件的廣義共軛余差法（P-GCR）

待解方程組Ax=b，解的初猜值x0，

余差ri， 下降方向向量pi，

預(yù)條件矩陣M-1，重啟間隔s，

系數(shù)αk,βkj

x0,r0=b-Ax0,z0=M-1r0,p0=z0

while not converged do

1.for k=0:s-1,do

2.αk=rTkApk/pTkATApk

3.xk+1=xk+αkpk

4.rk+1=rk-αkApk?zk+1=zk-αkM-1Apk

if converged，exit

5.βjk=-ATApj/pTjATApj(j=0,1,…,k)

6.pk+1=zk+1+

7.end for

8.r0=rs,p0=ps

end while

算法1中，A、M-1是n階方陣（n的大小取決于具體算例），x、r、p、z是長(zhǎng)度為n的向量。

算法1中第2步的pTkATApk和第5步的ATApj分別是一次長(zhǎng)向量點(diǎn)乘計(jì)算。長(zhǎng)向量的點(diǎn)乘需要進(jìn)行并行化計(jì)算，各個(gè)進(jìn)程計(jì)算結(jié)果進(jìn)行全局范圍的求和，由此要各產(chǎn)生一次全局通信，因此每s步迭代中（本文中取s=10），將進(jìn)行2×s次全局通信。并且通信間存在嚴(yán)格的依賴關(guān)系，無(wú)法通過(guò)合并而減少。

3.2 CA-P-GCR算法及推導(dǎo)

CA算法的基本思路：

（1）證明迭代過(guò)程中的長(zhǎng)向量都存在于同一個(gè)Krylov子空間中，因此可以用空間下一組坐標(biāo)（短向量）來(lái)代替長(zhǎng)向量。

（2）把短向量代入原始算法，用短向量的迭代取代相應(yīng)長(zhǎng)向量的迭代，以此實(shí)現(xiàn)迭代過(guò)程中的通信避免。

CA-P-GCR算法的推導(dǎo)：

定義Krylov子空間

Ks+1(A,v)=span{v,Av,A2v,…,Asv}（v為任意向量）

在算法1中，當(dāng)k=0時(shí)，由第3步得：

x1-x0∈K1(M-1A,p0)

由第4、6步得：

z1∈K2(M-1A,p0)，

p1∈K2(M-1A,p0)

同理，當(dāng)k=s時(shí)，有結(jié)論：

選取適當(dāng)向量建立基

其中，選取vp0=p0,vz0=z0。

vpi和vzi由下式產(chǎn)生：

vp(z)i=αM-1Avp(z)(i-1)+βvp(z)(i-1)+γvp(z)(i-2)

其中i≥2。

α、β、γ的值依所采用的不同基而定，較復(fù)雜的基（如切比雪夫基、牛頓基）之間具有更低的相關(guān)性，可以支持更高的重啟間隔s，但同時(shí)在計(jì)算α、β、γ時(shí)也會(huì)帶來(lái)一定的計(jì)算量。根據(jù)實(shí)際情況，在本例中，選取最簡(jiǎn)單的單項(xiàng)式基，即：α=1,β=0,γ=0

根據(jù)結(jié)論式（1），可以設(shè)：

其中，短向量mi,ni,li的長(zhǎng)度為2s+1。

將以上設(shè)定代入原算法中，代入第2步得：

代入第5步得：

設(shè)有Gram矩陣：

于是：

其中G、Gm均為2s+1階方陣。

由第3步得：

根據(jù)生成矩陣A的氣象學(xué)原理，當(dāng)p0≠z0時(shí)，V的各分量線性無(wú)關(guān)，方程有唯一零解；當(dāng)p0=z0時(shí)，V的秩為s+1，短向量的相應(yīng)坐標(biāo)也保持相等，方程相當(dāng)于一個(gè)s+1列方程，同樣只有唯一零解，因此有：

同理，由第6步得：

由第4步得：

設(shè)有換基矩陣T滿足M-1AV=VT，得：

其中，換基矩陣T在單項(xiàng)式基下為：

其廣義形式根據(jù)公式vp(z)i=αM-1Avp(z)(i-1)+βvp(z)(i-1)+γvp(z)(i-2)而定。

經(jīng)過(guò)這樣的替代，在迭代過(guò)程中，不再有長(zhǎng)向量出現(xiàn)（只在計(jì)算Gram矩陣過(guò)程中出現(xiàn)長(zhǎng)向量），于是，“通信避免的帶預(yù)條件的廣義共軛余差法”（Communication Avoiding Preconditioned Generalized Conjugate Residual method，CA-P-GCR，以下簡(jiǎn)稱CA算法）描述如下：

算法2通信避免的帶預(yù)條件的廣義共軛余差法

x0,r0=b-Ax0,z0=M-1r0,p0=z0

while not converged do

1.Calculate V=[p0,M-1Ap0,(M-1A)2p0,…,(M-1A)sp0,z0,M-1Az0,…,(M-1A)s-1z0]

2.Let G=VTATAV Gm=VTMTAV

3.m0=[0s+1,1,0s-1]Tn0=[1,02s]Tl0=[02s+1]T

4.for k=0:s-1，do

5.αk=mTkGmnk/nTkGnk

6.lk+1=lk+αknk

7.mk+1=mk-αkTnk

8.βjk=-Gnj/nTjGnj

9.nk+1=mk+1+

10.end for

11.zs=Vms,ps=Vns,xs=x0+Vls

12.z0=zs,p0=ps

end while

算法2中，0s表示s個(gè)實(shí)數(shù)0組成的行向量，只有在第2步與第3步之間進(jìn)行一次全局通信，即每s次迭代進(jìn)行一次全局通信。

為減少計(jì)算量，本文又對(duì)算法2中第11、12步修改如下：

11.xs=x0+Vls

12.r0=b-Axs,z0=M-1r0,p0=z0

修改前，第12步新生成的z0≠p0，其后的第1、2步的計(jì)算過(guò)程中要對(duì)z0、p0分別進(jìn)行稀疏矩陣向量乘、點(diǎn)乘。修改后，第12步新生成的z0=p0，其優(yōu)點(diǎn)是大幅減少了后續(xù)第1、2步的計(jì)算量。相應(yīng)的代價(jià)是總迭代次數(shù)會(huì)有小幅增加，因?yàn)檫@樣修改之后的算法和原算法在數(shù)學(xué)上已經(jīng)不再等價(jià)。但是由于CA算法的迭代次數(shù)只能是s的整數(shù)倍，實(shí)際運(yùn)行后發(fā)現(xiàn)，實(shí)際的迭代次數(shù)并沒(méi)有因?yàn)橐陨闲薷亩黾印?/p>

3.3 解的唯一性與正確性

待解方程Ax=b的解的唯一性由A、b的屬性決定，而A、b的生成依賴于相關(guān)的氣象學(xué)原理。根據(jù)相關(guān)氣象學(xué)原理可以確定，待解方程的解存在且唯一。

新算法與舊算法求得的方程解并不完全一致，其原因至少包括以下2條：

（1）兩種算法迭代次數(shù)不同。

（2）算法2中對(duì)第11、12步的修改。

其中，原因（1）在所有CA類算法中普遍存在。但是模式動(dòng)力框架對(duì)方程求解器模塊的要求，并不是不同算法解的一致，而是僅僅要求收斂。后續(xù)的動(dòng)力框架部分和物理過(guò)程部分的計(jì)算，也都僅以方程求解器的收斂為前提。

3.4 新舊算法計(jì)算量對(duì)比

對(duì)比原算法和CA-P-GCR算法，計(jì)算量對(duì)比如表1。

表1 新舊算法的計(jì)算量

對(duì)比發(fā)現(xiàn)：每s步迭代（即每個(gè)restart周期），CA-P-GCR中全局通信的次數(shù)減少為原算法的1/(2×s)。稀疏矩陣向量乘減少了約一半，點(diǎn)乘計(jì)算量增加s次。按照以往測(cè)試結(jié)果，稀疏矩陣向量乘的計(jì)算量遠(yuǎn)高于點(diǎn)乘。因此預(yù)測(cè)CA算法將在計(jì)算、通信兩方面同時(shí)優(yōu)于原算法。同時(shí)，相比原算法，計(jì)算部分更加集中，存在進(jìn)行訪存優(yōu)化、提高計(jì)算訪存比的潛力。

4 數(shù)值實(shí)驗(yàn)

4.1 算例信息

本測(cè)試的測(cè)試對(duì)象是中國(guó)氣象局的新一代“全球區(qū)域一體化數(shù)值預(yù)報(bào)模式（GRAPES）”，其動(dòng)力框架的核心是求解亥姆霍茲方程Ax=b，本測(cè)試共運(yùn)行288步，每步進(jìn)行1次方程求解，每次求解過(guò)程中，A、b不變，x初猜值繼承上一次的最終結(jié)果（第一次初猜值取0向量）。本測(cè)試采用0.05°水平分辨率的全球算例，垂直層數(shù)60層，時(shí)間步長(zhǎng)150 s，物理過(guò)程關(guān)閉。A是1.56×109階方陣，x、b是長(zhǎng)度為1.56×109的向量。因?yàn)锳、x、b以及迭代過(guò)程中出現(xiàn)的z、p、r規(guī)模較大，所以計(jì)算時(shí)按節(jié)點(diǎn)分別存儲(chǔ)。A是稀疏矩陣，每一行中有19個(gè)非0元。

4.2 計(jì)算集群信息

本測(cè)試所用計(jì)算集群是中國(guó)氣象局2018年部署的“曙光-派”計(jì)算集群，采用Intel的SKL Gold 6142處理器，主頻2.6 GHz，32核/節(jié)點(diǎn)，每個(gè)核心運(yùn)行1個(gè)進(jìn)程，純MPI并行。每節(jié)點(diǎn)內(nèi)存為192 GB的DDR4內(nèi)存，操作系統(tǒng)RHEL 7.4，編譯器Intel PSXE 2017u2。

5 實(shí)驗(yàn)結(jié)果

5.1 收斂速率對(duì)比

測(cè)試中，每種算法各進(jìn)行方程求解288次，每次求解的迭代次數(shù)各不相同。經(jīng)統(tǒng)計(jì)得：CA算法平均迭代次數(shù)為34次，原算法平均迭代次數(shù)為28次?，F(xiàn)以4 096進(jìn)程下一次典型求解過(guò)程為例，如圖1所示。

圖1 新舊算法收斂速率

說(shuō)明：本文以r0全體元素的平方和作為殘差（由于三維空間格點(diǎn)總數(shù)不變，因此沒(méi)有進(jìn)行平均、開(kāi)方運(yùn)算），收斂閾值定為1.0×10-9，圖1中縱坐標(biāo)為殘差的常用對(duì)數(shù)。由圖1可知，本次求解過(guò)程中，隨著迭代次數(shù)的增加，CA算法的收斂速率略慢于原算法。同時(shí)，CA算法的迭代次數(shù)為不連續(xù)的、10的整數(shù)倍。原算法達(dá)到收斂閾值的迭代次數(shù)是33次，CA算法達(dá)到收斂閾值的次數(shù)是40次。

在整個(gè)測(cè)試中，CA算法的平均迭代次數(shù)高于原算法的原因主要有以下3條：

（1）CA算法的迭代次數(shù)只能是s的整數(shù)倍（本文中s全部取10）。

（2）所有CA算法由于基的條件數(shù)的限制，都存在最高收斂精度降低，收斂速率小幅變慢的問(wèn)題。Carson于2014年針對(duì)各種“通信避免算法”中出現(xiàn)的這類問(wèn)題進(jìn)行了詳細(xì)分析，并給出了“余差替代”解決方案[9]。本文所提出的算法也存在相同的問(wèn)題，但由于“余差替代”方案開(kāi)銷較大，因此沒(méi)有采納。

（3）前文中修改第12步，改變生成新的z0、p0的方法，使得z0、p0相等。這就改變了基的條件數(shù)，進(jìn)一步減慢了收斂速率。

由于要兼顧通信量和計(jì)算量，s的選取受到一定限制，使得以上（2）、（3）兩條的效果被（1）所覆蓋，因此（2）、（3）兩條并沒(méi)有實(shí)質(zhì)上增加總的迭代次數(shù)[10]。

5.2 總時(shí)間對(duì)比

本測(cè)試并行規(guī)模從2 048進(jìn)程到16 384進(jìn)程，GRAPES模式的整體運(yùn)行時(shí)間（圖2）和亥姆霍茲方程求解器的運(yùn)行時(shí)間及加速比（圖3）如下所示。

圖2 新舊算法的模式總運(yùn)行用時(shí)及加速比

圖3 新舊算法的求解器運(yùn)行用時(shí)及加速比

對(duì)比圖2和圖3可知，方程求解器的時(shí)間消耗占據(jù)了整個(gè)數(shù)值模式的重要部分（從42%到65%）。求解器之外的部分并沒(méi)有做修改，因此差別不明顯。求解器算法的優(yōu)化，導(dǎo)致了模式整體運(yùn)行速率最高1.64倍的加速比（相對(duì)于同規(guī)模下的原算法）[11]。

由表1可知，CA算法在計(jì)算量、通信量上都優(yōu)于原算法，從圖3也可知從2 048進(jìn)程到16 384進(jìn)程上，CA算法的用時(shí)都少于原算法，并且這種優(yōu)勢(shì)隨著并行規(guī)模擴(kuò)展而擴(kuò)大。從圖3的加速比可以看出，當(dāng)規(guī)模擴(kuò)展到16 384進(jìn)程時(shí)，CA算法的加速比達(dá)到了原算法的3倍[12]。對(duì)比新舊算法求解器部分的可擴(kuò)展性，兩種算法在規(guī)模較小時(shí)都表現(xiàn)出良好的可擴(kuò)展性，當(dāng)規(guī)模較大時(shí)，CA算法仍然保持可擴(kuò)展性，而原算法隨規(guī)模擴(kuò)展而速度降低。

5.3 主要計(jì)算、通信部分時(shí)間對(duì)比

并行規(guī)模較小時(shí)，計(jì)算開(kāi)銷（包括稀疏矩陣向量乘、點(diǎn)乘）占據(jù)主要部分，其用時(shí)遠(yuǎn)高于通信。并且隨著規(guī)模的擴(kuò)展，計(jì)算部分的加速呈現(xiàn)“超線性”[13]，例如：4 096進(jìn)程的并行規(guī)模是2 048進(jìn)程的2倍，但稀疏矩陣向量乘用時(shí)僅為2 048進(jìn)程的1/4～1/3。推測(cè)其主要原因是規(guī)模較小情況下，內(nèi)存壓力較大，導(dǎo)致訪存命中率偏低。隨著并行規(guī)模擴(kuò)大，通信用時(shí)逐漸占據(jù)了主要部分，CA算法的優(yōu)勢(shì)主要體現(xiàn)在通信上，16 384進(jìn)程下，原算法的通信開(kāi)銷已經(jīng)遠(yuǎn)遠(yuǎn)超過(guò)了計(jì)算開(kāi)銷。

另外，對(duì)比圖2、圖4，隨著并行規(guī)模的擴(kuò)展，原算法的全局通信用時(shí)最高占據(jù)了模式整體運(yùn)行時(shí)間的50%，由此說(shuō)明了通信避免的必要性。

圖4 新舊算法的主要計(jì)算及通信用時(shí)

6 結(jié)論

GRAPES模式使用原算法的亥姆霍茲方程求解器受限于全局通信，其擴(kuò)展能力止步于8 192進(jìn)程。本文通過(guò)通信避免算法改造，以小幅降低收斂速率為代價(jià)，同時(shí)改善了求解器的運(yùn)行速率和擴(kuò)展能力，在16 384規(guī)模平臺(tái)上的實(shí)驗(yàn)顯示出3倍于原算法的加速效果[14]。

綜合以上結(jié)果可知：CA算法可以極大地改善方程求解器的可擴(kuò)展性，并且減少運(yùn)算、通信的時(shí)間開(kāi)銷，在各種計(jì)算規(guī)模下都有著相比于原算法明顯的優(yōu)勢(shì)[15]。在較小規(guī)模下的優(yōu)勢(shì)主要來(lái)源于計(jì)算的減少，在較大規(guī)模下的優(yōu)勢(shì)主要來(lái)源于全局通信的減少。同時(shí)，會(huì)導(dǎo)致收斂速率的小幅降低。