廣義共軛余差法的通信避免算法

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

1.中國氣象科學研究院，北京100081

2.英特爾中國有限公司，北京100081

1 引言

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

一般來說，算法的開銷可以歸結為兩部分：數據的計算和數據的移動。其中數據的移動包括節點內部的移動和節點間通過網絡的移動。要減少這種密集的數據移動帶來的開銷，加州大學伯克利分校的Demmel教授最早于2007年提出了通信避免（Communication Avoiding，CA）算法[3]。CA算法的思路是：構造能包含迭代過程中的所有長向量的Krylov子空間，利用短向量的迭代替代長向量的迭代，進而避免在迭代過程中進行通信。該思路具體應用于對不同的算法的改造，會產生相應的不同的CA算法。目前針對GCR算法，相關工作仍是空白。

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

2 相關工作

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

3 算法描述

3.1 原始算法

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

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

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

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

預條件矩陣M-1，重啟間隔s，

系數α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是長度為n的向量。

算法1中第2步的pTkATApk和第5步的ATApj分別是一次長向量點乘計算。長向量的點乘需要進行并行化計算，各個進程計算結果進行全局范圍的求和，由此要各產生一次全局通信，因此每s步迭代中（本文中取s=10），將進行2×s次全局通信。并且通信間存在嚴格的依賴關系，無法通過合并而減少。

3.2 CA-P-GCR算法及推導

CA算法的基本思路：

（1）證明迭代過程中的長向量都存在于同一個Krylov子空間中，因此可以用空間下一組坐標（短向量）來代替長向量。

（2）把短向量代入原始算法，用短向量的迭代取代相應長向量的迭代，以此實現迭代過程中的通信避免。

CA-P-GCR算法的推導：

定義Krylov子空間

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

在算法1中，當k=0時，由第3步得：

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

由第4、6步得：

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

p1∈K2(M-1A,p0)

同理，當k=s時，有結論：

選取適當向量建立基

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

vpi和vzi由下式產生：

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

其中i≥2。

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

根據結論式（1），可以設：

其中，短向量mi,ni,li的長度為2s+1。

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

代入第5步得：

設有Gram矩陣：

于是：

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

由第3步得：

根據生成矩陣A的氣象學原理，當p0≠z0時，V的各分量線性無關，方程有唯一零解；當p0=z0時，V的秩為s+1，短向量的相應坐標也保持相等，方程相當于一個s+1列方程，同樣只有唯一零解，因此有：

同理，由第6步得：

由第4步得：

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

其中，換基矩陣T在單項式基下為：

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

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

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

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個實數0組成的行向量，只有在第2步與第3步之間進行一次全局通信，即每s次迭代進行一次全局通信。

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

11.xs=x0+Vls

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

修改前，第12步新生成的z0≠p0，其后的第1、2步的計算過程中要對z0、p0分別進行稀疏矩陣向量乘、點乘。修改后，第12步新生成的z0=p0，其優點是大幅減少了后續第1、2步的計算量。相應的代價是總迭代次數會有小幅增加，因為這樣修改之后的算法和原算法在數學上已經不再等價。但是由于CA算法的迭代次數只能是s的整數倍，實際運行后發現，實際的迭代次數并沒有因為以上修改而增加。

3.3 解的唯一性與正確性

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

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

（1）兩種算法迭代次數不同。

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

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

3.4 新舊算法計算量對比

對比原算法和CA-P-GCR算法，計算量對比如表1。

表1 新舊算法的計算量

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

4 數值實驗

4.1 算例信息

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

4.2 計算集群信息

本測試所用計算集群是中國氣象局2018年部署的“曙光-派”計算集群，采用Intel的SKL Gold 6142處理器，主頻2.6 GHz，32核/節點，每個核心運行1個進程，純MPI并行。每節點內存為192 GB的DDR4內存，操作系統RHEL 7.4，編譯器Intel PSXE 2017u2。

5 實驗結果

5.1 收斂速率對比

測試中，每種算法各進行方程求解288次，每次求解的迭代次數各不相同。經統計得：CA算法平均迭代次數為34次，原算法平均迭代次數為28次。現以4 096進程下一次典型求解過程為例，如圖1所示。

圖1 新舊算法收斂速率

說明：本文以r0全體元素的平方和作為殘差（由于三維空間格點總數不變，因此沒有進行平均、開方運算），收斂閾值定為1.0×10-9，圖1中縱坐標為殘差的常用對數。由圖1可知，本次求解過程中，隨著迭代次數的增加，CA算法的收斂速率略慢于原算法。同時，CA算法的迭代次數為不連續的、10的整數倍。原算法達到收斂閾值的迭代次數是33次，CA算法達到收斂閾值的次數是40次。

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

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

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

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

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

5.2 總時間對比

本測試并行規模從2 048進程到16 384進程，GRAPES模式的整體運行時間（圖2）和亥姆霍茲方程求解器的運行時間及加速比（圖3）如下所示。

圖2 新舊算法的模式總運行用時及加速比

圖3 新舊算法的求解器運行用時及加速比

對比圖2和圖3可知，方程求解器的時間消耗占據了整個數值模式的重要部分（從42%到65%）。求解器之外的部分并沒有做修改，因此差別不明顯。求解器算法的優化，導致了模式整體運行速率最高1.64倍的加速比（相對于同規模下的原算法）[11]。

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

5.3 主要計算、通信部分時間對比

并行規模較小時，計算開銷（包括稀疏矩陣向量乘、點乘）占據主要部分，其用時遠高于通信。并且隨著規模的擴展，計算部分的加速呈現“超線性”[13]，例如：4 096進程的并行規模是2 048進程的2倍，但稀疏矩陣向量乘用時僅為2 048進程的1/4～1/3。推測其主要原因是規模較小情況下，內存壓力較大，導致訪存命中率偏低。隨著并行規模擴大，通信用時逐漸占據了主要部分，CA算法的優勢主要體現在通信上，16 384進程下，原算法的通信開銷已經遠遠超過了計算開銷。

另外，對比圖2、圖4，隨著并行規模的擴展，原算法的全局通信用時最高占據了模式整體運行時間的50%，由此說明了通信避免的必要性。

圖4 新舊算法的主要計算及通信用時

6 結論

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

綜合以上結果可知：CA算法可以極大地改善方程求解器的可擴展性，并且減少運算、通信的時間開銷，在各種計算規模下都有著相比于原算法明顯的優勢[15]。在較小規模下的優勢主要來源于計算的減少，在較大規模下的優勢主要來源于全局通信的減少。同時，會導致收斂速率的小幅降低。