GROMACS 2020在ROCm平臺上的移植與優(yōu)化*

張馭洲，曹武迪，卜景德，譚光明，吉 青

(1.中國科學院理論物理研究所理論物理先進計算聯合實驗室，北京 100190； 2.中國科學院計算技術研究所計算機體系結構國家重點實驗室， 北京 100190)

1 引言

分子動力學MD(Molecular Dynamics)[1]模擬是利用計算機模擬原子或分子運動的方法，通過模擬微觀層面大量原子或分子的運動來研究物質的結構和性質。由于其可以方便地模擬實驗方法難以測定的微觀結構和物理過程而廣泛應用于生物、化學、醫(yī)藥和材料等研究領域。經典MD模擬方法建立在牛頓力學基礎之上，通過計算模擬體系中大量原子或分子之間的相互作用，最終得到整個體系的宏觀統(tǒng)計信息。當前具有實際應用意義的MD模擬通常需要對包含至少數萬、數十萬粒子的體系進行百萬乃至上億時間步的模擬計算，在空間和時間尺度上都對計算能力有著極高的要求，因此，一直以來，MD都是高性能計算HPC(High Performance Computing)領域的一個重要應用方向。

GROMACS[2]是一個應用廣泛的開源MD模擬軟件。在眾多的開源MD軟件之中，GROMACS以高性能和高效率著稱。當前世界上最大的分布式計算項目Folding@home[3]使用的計算引擎就是GROMACS。GROMACS擁有高度優(yōu)化的CPU計算代碼，支持通過MPI+OpenMP進行大規(guī)模并行計算，還通過CUDA支持NVIDIA GPU，并通過OpenCL支持其他加速設備，但OpenCL版本對新功能的支持一般落后于CUDA。當前最新的GROMACS 2020系列已經通過CUDA將MD模擬中計算量最大的3種作用力類型——短程非成鍵作用力、長程作用力和成鍵作用力均實現了GPU加速計算，性能得到了進一步提高[4]。而OpenCL版本目前尚不支持成鍵力在GPU上的計算。

ROCm(Radeon Open Compute)[5]是AMD推出的一個開源高性能異構計算平臺，提供了包括GPU驅動、編譯器、運行時、數學庫、性能分析與調試工具在內的一整套編程開發(fā)和程序運行環(huán)境。ROCm提供了一個名為HIP(Heterogeneous-computing Interface for Portability)的編程模型及語言[6]，使用HIP編寫的應用程序可以通過ROCm平臺在AMD GPU上運行。HIP的應用層API與CUDA相似，可以較為方便地將現有CUDA程序轉碼為HIP程序，從而將只能在NVIDIA GPU上運行的應用移植到AMD GPU上。根據媒體披露，美國的3臺E級高性能計算機中，有2臺都將采用AMD GPU。

作為開源平臺，ROCm擁有比閉源的CUDA更好的靈活性，但其生態(tài)的建設離不開豐富的應用的支持。本文通過將最新的GROMACS 2020系列轉碼為HIP代碼，實現了向ROCm平臺的移植，并通過對代碼的一系列優(yōu)化，在目標算例上獲得了相對初始轉碼版本約2.8倍的加速比。據我們所知，這是目前世界范圍內第一個實現最新版GROMACS完整功能的HIP版本。

本文主要有以下貢獻：

(1)在ROCm平臺首次實現了GROMACS 2020系列的完整功能，一方面為ROCm平臺的生態(tài)系統(tǒng)添加了GROMACS這一應用廣泛的HPC軟件，另一方面移植的過程也可為其他應用的移植提供參考。

(2)以有代表性的算例為目標，在單結點內分析了其計算熱點與瓶頸，進行了針對性的性能優(yōu)化，獲得了相對初始轉碼版本約2.8倍的加速比，其優(yōu)化思路與方法也可以為其他GPU應用的優(yōu)化工作提供借鑒；通過多結點擴展性測試，分析了擴展性方面的瓶頸，為后續(xù)優(yōu)化打下了基礎。

2 程序移植

HIP是ROCm 運行時環(huán)境之上的一層很薄的封裝，它提供了一系列代碼自動檢查與轉換工具，可以自動檢查源文件中的CUDA API并轉換為相應的HIP API，因此簡單程序的移植已經十分方便。然而，大型的應用程序文件數量多，結構復雜，不僅代碼文件需要轉換，編譯構建工具也需要修改，因此往往需要更多人工調試才能成功移植。

GROMACS經歷了二十多年的不斷開發(fā)與迭代，現在已經成為了一個龐大而復雜的應用軟件，代碼量已達百萬行量級[2]。GROMACS 2020.2版本源代碼基本組織結構如圖1所示。主要的源碼文件都位于src/gromacs目錄下。GROMACS使用CMake[7]進行編譯構建，頂層目錄的cmake目錄下是配置各種編譯參數的cmake文件，在各級源碼目錄下都包括一個CMakeLists.txt文件，用于管理本級目錄下源碼文件的編譯。通過使用HIP提供的工具完成自動代碼轉換之后，本文對代碼和構建系統(tǒng)進行了調試修改，完成了編譯，構建的GROMACS可以通過所有自帶的單元測試。

Figure 1 Source code organization of GROMACS 2020.2圖1 GROMACS 2020.2 源代碼基本組織結構

3 性能優(yōu)化

3.1 測試算例與基準

為保持與先前工作的一致性[8]，本文使用的算例是 [C16MP]+[Br3]-。首先驗證HIP版GROMACS的計算正確性，然后作為性能基準，進行后續(xù)的代碼優(yōu)化。該算例模擬體系為離子液體體系，由1 440對離子構成，原子數為95 040，有大量成鍵相互作用(bonded interaction)和非成鍵相互作用nb(non-bonded interaction)。測試使用的軟硬件版本及型號列于表1。分別使用GROMACS 2020.2原版CPU程序、原版OpenCL程序和本文移植的HIP版本程序，采用相同的模擬參數(版本功能特性相關參數除外)將測試算例運行5萬時間步，運行的命令行參數列于表2。以原版CPU程序5萬步平均計算結果為基準，OpenCL版和HIP版的相對誤差結果都在同一數量級，并且都在MD允許的波動范圍內，證明了HIP版計算結果的正確性；最后1萬步統(tǒng)計的性能數據如表3所示。

Table 1 Software hardware and versions and models used in the test表1 本文測試使用的軟硬件版本及型號

Table 2 Command line parameters used in the test表2 測試中使用的命令行參數

Table 3 Performance comparison of three versions of GROMACS on 10 000 steps running表3 3個版本GROMACS運行10 000步的性能數據

可見原版的OpenCL版本比CPU版本約有67.9%的加速，而移植的HIP版本性能甚至略低于CPU版本。

3.2 性能分析與瓶頸定位

MD的計算量主要來源于粒子之間的短程非成鍵相互作用(short-ranged non-bonded interaction)、長程相互作用(通過PME(Particle-Mesh-Ewald)方法計算)和成鍵相互作用。本文移植的HIP版GROMACS已經將這3種作用力的計算全部都在GPU上實現了，因此，需要一個GPU核函數性能分析工具來考察程序的性能瓶頸。ROCm提供了rocprof[9]工具來完成這一工作。rocprof可以在runtime層抓取GPU的運行數據，獲得包括核函數的發(fā)起與結束時間、內存拷貝的起始、結束以及數據量等信息，還可以結合Chrome瀏覽器的tracing工具實現可視化。

Figure 2 Time lines of kernel of the first ported HIP code圖2 HIP初始版本的核函數時間線

圖2為HIP初始版本的核函數時間線。圖2中顯示的核函數執(zhí)行順序與GROMACS計算流程一致，但是每個函數的執(zhí)行時間都異常長，因此本文后續(xù)著重對每個核函數進行了優(yōu)化。

3.3 優(yōu)化方法與結果

3.3.1 bonded核函數優(yōu)化

bonded核函數是一個函數模板，在一個50時間步的運行中，其不同的實例調用次數、平均運行時間及總時間如表4所示。

Table 4 Performance analysis results of the bonded kernel表4 bonded 核函數性能分析結果

不同的模板參數下，bonded核函數的運行時間差異較大。考察其代碼，發(fā)現不同模板參數的實例的差異主要在于是否計算位力(virial，作用于粒子上的合力與粒子矢徑的標積)和能量。位力和能量的計算需要將體系中所有粒子的貢獻都累加起來，GROMACS原版的CUDA代碼直接使用atomicAdd函數將每個線程的計算結果累加到GPU全局內存中的變量，這一操作在MI50 GPU上導致了嚴重的性能下降。MI50 GPU硬件支持32位整型的原子加操作，但是浮點型原子加操作是通過一個CAS(Compare-And-Swap)循環(huán)實現的，性能較差，大量對全局內存地址的原子加操作會嚴重降低程序性能。針對這種情況，本文將bonded核函數中對位力和能量的歸約進行了修改，首先在block內部利用共享內存進行歸約，然后再使用原子加在block之間進行歸約。利用共享內存上較快的原子加操作代替了大量直接操作全局內存地址的原子加操作，相關核函數性能有了大幅提升。

此外，本文還修改了bonded核函數的發(fā)起參數，將原來256線程的block尺寸縮減為64線程，性能得到了進一步提高。MD中成鍵力的計算量與非成鍵力相比要小很多，并且有較多邏輯判斷，嚴格意義上并不適合GPU計算。GROMACS原版的CUDA代碼中，bonded核函數有一個包括8種成鍵力的switch語句。當某些warp發(fā)生分支，整個block的計算都將被減慢。將block尺寸減小到ROCm平臺上warp(wavefront)的尺寸，就完全消除了正常warp對發(fā)生分支的warp的等待，從而進一步縮短了bonded核函數的平均執(zhí)行時間。

圖3為兩步優(yōu)化前后bonded核函數的性能。由于原來原子加操作只有在計算位力和能量時才會執(zhí)行，因此第1步減少原子加操作的優(yōu)化大幅減少了模板參數為〈true,true〉的核函數實例的運行時間，而其他2種實例性能不變。通過減小block尺寸的進一步優(yōu)化，bonded核函數的3種實例性能都得到大幅提升，其中后2種實例執(zhí)行時間甚至降為0，是因為單個核函數的執(zhí)行時間已經小于性能分析工具所能抓取的最小時間間隔。

Figure 3 Average execution time of bonded kernels before and after optimization圖3 優(yōu)化前后bonded核函數的平均執(zhí)行時間

3.3.2 PME核函數優(yōu)化

PME部分的2個主要核函數pme_solve_kernel和pme_spline_and_spread_kernel中也使用了直接操作GPU全局內存的浮點型原子加操作，本文也對這2個函數進行了優(yōu)化嘗試。

(1)pme_solve_kernel。

當需要計算位力和能量時，pme_solve_kernel核函數的每個線程需要計算6個位力分量和1個能量變量，這7個量需要全局累加。原版代碼首先使用shuffle指令和共享內存，將每個線程所計算的7個變量的結果歸約到線程所屬block的前7個線程，然后每個block的前7個線程使用浮點型原子加操作將結果歸約到全局內存。盡管已經通過shuffle指令和共享內存歸約了block內的結果，但由于block數量眾多，所有的block都對全局內存中同樣7個地址進行原子加操作，仍然嚴重影響了性能。

Figure 4 Schematic diagram of pme_solve_kernel optimization 圖4 pme_solve_kernel優(yōu)化示意圖

為了避免使用性能不佳的原子加操作，本文在原來保存7個位力和能量變量的GPU全局內存地址之后再額外申請一段內存，其長度為pme_solve核函數的block數乘以7個float變量大小。pme_solve_kernel在計算得到7個位力和能量變量之后不再向全局內存歸約，而是直接寫入額外申請的全局內存中，按照分量的種類寫入其對應的內存塊，每種分量數量等于pme_solve_kernel的block數，總共7種分量。在全部的pme_solve_kernel核函數完成計算和全局內存寫入之后，使用另外一個單獨的核函數讀取pme_solve_kernel核函數寫入全局內存的結果，使用shuffle指令和共享內存進行block內的歸約。該單獨歸約的核函數的block尺寸設置為MI50 GPU允許的最大blockSize 1 024，并且只發(fā)起一個block，這樣block內的歸約結果就是最終結果，從而完全避免了對全局內存原子加操作的使用。圖4為pme_solve_kernel優(yōu)化示意圖。

Figure 5 Average execution time of the pme_solve_kernel before and after optimization圖5 pme_solve_kernel核函數 優(yōu)化前后核函數平均執(zhí)行時間

圖5為優(yōu)化前后pme_solve_kernel不同實例的平均執(zhí)行時間?？梢妼υ蛹硬僮鞯膬?yōu)化大幅提升了計算位力和能量的核函數實例的性能，其執(zhí)行時間已經接近不計算位力和能量的核函數實例。當然，為了替代原子加，本文引入單獨歸約核函數，該核函數平均執(zhí)行時間約15 μs，因此總體上仍然提升了代碼性能。

(2)pme_spline_and_spread_kernel。

pme_spline_and_spread_kernel核函數的一個功能是將每個粒子的電荷分散到周圍n3個FFT(Fast Fourier Transform)網格點上，n為B樣條插值的階數。該操作在block內的原子加操作寫入同一全局內存地址的比例很少，而不同的block之間則有較多訪問同一全局內存地址的原子加操作，因此無法通過先行歸約block內的線程來減少原子加操作的數量。而對于在pme_solve_kernel中使用的方法，即通過單獨的歸約核函數來避免原子加操作，這里也不再有效。這是因為，從歸約的目標地址和每個目標地址上所需歸約的元素的數量來看，pme_solve_kernel的原子加操作是“密集型”：歸約目標為7個全局內存地址，每個地址需要累加至少上千個元素；而pme_spline_and_spread_kernel的原子加操作是“稀疏型”：歸約目標為整個體系FFT網格的所有格點，每個格點需要累加的元素數目不定，而是根據該點周圍的帶電粒子密度而變化的，一般多則幾十個，少的甚至為零，而FFT網格的數目卻很大，以本算例為例，FFT網格尺寸為72 × 48 × 192，共663 552個格點。pme_solve_kernel采用的單獨歸約核函數方法，在這種情況下將變得極為低效：首先，需要額外申請的內存空間很大，對于本例來說，即使假設每個目標地址最多只需要累加10個元素，總共也需要6 635 520個float數，超過25 MB的全局內存；而且由于每個時間步每個目標地址所需累加的元素個數可能改變，每一步計算都需要將這些內存全部清零。其次，對于這樣大量的數據，采用單獨的歸約核函數從全局內存讀入再歸約，也非常耗時。本文實現了該方法并進行了測試，結果表明，在pme_spline_and_spread_kernel核函數中使用單獨的歸約核函數后，性能比原來直接使用原子加更差。因此，本文對于該核函數沒有進行進一步的優(yōu)化，仍然使用原版直接原子加操作的方案。這一案例也體現了從硬件上直接加速原子加操作的必要性。

3.3.3 nb核函數優(yōu)化

短程非成鍵相互作用(nb)是MD中計算量最大的部分，并且非常適合GPU計算，因此也是GROMACS最早實現GPU計算的部分。GROMACS為nb設計了一個精巧的算法[10]，在CPU和GPU上都獲得了良好的性能。隨著GROMACS支持的模擬種類的不斷豐富，nb核函數已經成為了一個包含數十種版本的函數模板。本文針對MI50 GPU的硬件特性，對[C16MP]+[Br3]-算例所使用的nb核函數實例代碼進行了2個方面的優(yōu)化，獲得了一定的性能提升。

(1)減少原子加操作數量。

nb核函數同樣使用了訪問全局內存地址的浮點型原子加操作，分別用于歸約每個粒子受到的鄰居粒子的作用力和偏移力。原版CUDA代碼中，對于大部分CUDA設備，nb核函數的每個block都是64線程(計算能力為3.7時為128線程)，占用2個warp，因此在歸約時也是按照2個warp進行設計。在MI50 GPU上，nb核函數剛好只占用1個warp，因此我們把相關的歸約均改為warp內的shuffle操作，將訪問全局內存的浮點型原子加操作數量減少了一半，提高了該核函數的性能。

(2)減少寄存器使用量。

盡管GPU被設計為用于大規(guī)模的并行計算，但是其并行性也是受硬件資源數量限制的，主要的限制因素是核函數中共享內存以及寄存器文件的使用量。nb是一個規(guī)模較大的核函數，使用了較多的寄存器變量。本文將編譯該核函數得到的目標文件反匯編后，發(fā)現其占用的寄存器數量達到97個，在ROCm平臺，這一寄存器使用量導致MI50 GPU上的每個單指令多數據SIMD(Single Instruction Multiple Data)計算組只能并發(fā)執(zhí)行2個warp。較低的并發(fā)數使得程序運行過程中某些指令(例如原子加操作)的延遲難以得到掩蓋。

nb核函數雖然使用的寄存器數量較多，但共享內存的用量卻相對較少，如能將部分寄存器變量轉為使用共享內存，有望提高并發(fā)數。然而，共享內存的讀寫速率仍然是遠低于寄存器變量的，實際驗證中發(fā)現，在循環(huán)中有寫操作的寄存器變量轉移到共享內存后核函數性能反而下降。經過對代碼的分析發(fā)現，nb_sci是一個包含4個int分量的結構體變量，它在所有循環(huán)開始之前賦值1次，后續(xù)的循環(huán)之中只有讀操作，沒有寫操作。將nb_sci變量轉移到共享內存之后，nb核函數寄存器使用量降低為81個，這樣SIMD并發(fā)的warp數增長為3，經過實際測試，執(zhí)行次數占絕對多數的nb核函數實例nbnxn_kernel_ElecEw_VdwLJCombLB_F_cuda性能有約7%的提升，因此，在另外2個nb核函數性能基本不變甚至略有下降的情況下，程序的總體性能仍有提升。圖6為優(yōu)化前后nb核函數實例的平均執(zhí)行時間。

Figure 6 Average execution time of the nb kernels before and after optimization圖6 優(yōu)化前后nb核函數的平均執(zhí)行時間

3.4 其他算例性能測試

通過對MD模擬中計算量最大的3部分(即短程非成鍵相互作用、長程相互作用以及成鍵相互作用)的GPU核函數nb、PME和bonded的優(yōu)化，[C16MP]+[Br3]-算例的性能從最初的13.109 ns/d增長為36.980 ns/d，提升182.1%。

表5列出了實施一系列優(yōu)化措施后算例的整體運行時間，時間統(tǒng)計方法與表3數據統(tǒng)計方法一致，并計算了各個版本以無優(yōu)化HIP版本為基準的加速比。為進行對比，表5中也加入了表3中CPU和OpenCL程序的運行時間?？梢娊涍^優(yōu)化之后，HIP版本性能已經超過原版CPU和OpenCL版本的。

為驗證優(yōu)化的通用性，本文還對另外2個典型算例——Lysozyme[11]和waterbox[12]進行了性能對比測試。Lysozyme是對水盒子中的溶菌酶蛋白質的模擬，體系原子數為71 844，相對[C16MP]+[Br3]-，其成鍵力計算要少得多。waterbox是僅包含溶劑水分子的體系，本文使用的是包含64萬個水分子、共192 000個原子的體系，規(guī)模較大，但是成鍵力全部以“約束”(constraint)的方式在CPU上計算。表6列出了Lysozyme和waterbox算例的運行性能數據，可見優(yōu)化的HIP版本相對優(yōu)化之前以及原版的CPU和OpenCL版本均有一定的性能提升。

Table 5 Performance comparison among different optimization methods表5 不同優(yōu)化方法性能對比

Table 6 Performance comparison among different GROMACS versions with Lysozyme and waterbox benchmarks表6 算例Lysozyme和waterbox在不同版本GROMACS下的性能對比

進一步對比這2個算例的性能數據，運行Lysozyme算例時OpenCL版比CPU版的性能高71.6%，而運行waterbox算例時OpenCL版比CPU版的性能高97.4%，可見GROMACS作為一個通用MD模擬軟件，支持的模擬種類眾多，由于不同類型的算例實際運行的代碼不同，即使是官方原版，計算不同算例的相對性能也有差異。本文以[C16MP]+[Br3]-這一算例為目標進行的一系列優(yōu)化，對Lysozyme和waterbox算例也有一定加速效果，證明了優(yōu)化措施的通用性；但另一方面，在Lysozyme和waterbox算例上的加速效果不如[C16MP]+[Br3]-，這也體現了通用代碼的復雜性。若要獲得最佳的性能，則需要對目標算例進行針對性的分析與優(yōu)化。

3.5 擴展性測試

為了從更多角度測評優(yōu)化效果，本文還使用優(yōu)化前后的HIP代碼分別進行了擴展性測試，包括弱擴展性測試和強擴展性測試。由于PME算法本身在計算過程中要求全局通信，因此GROMACS 2020版本只支持單GPU卡計算PME，即使使用多結點運行，也只能指定使用一個進程單獨計算PME，通常這會對程序的擴展性造成不利影響。使用CPU計算PME則沒有這一限制，因此，在擴展性測試部分，各個規(guī)模下均使用CPU計算PME，bonded和nb部分仍然使用GPU進行計算。

弱擴展性測試將[C16MP]+[Br3]-這一算例規(guī)模分別擴展為原來的1倍、2倍、4倍和8倍，分別使用1,2,4和8個結點進行計算。強擴展性測試則將[C16MP]+[Br3]-這一算例擴展為原來的8倍，分別使用1，2，4和8個結點進行計算。擴展性測試中時間統(tǒng)計方法與表3數據統(tǒng)計方法一致，結果分別列于表7和表8。

Table 7 Test results of weak scalability 表7 弱擴展性測試結果

Table 8 Test results of strong scalability 表8 強擴展性測試結果

由于PME部分采用CPU計算，所以損失了該部分的優(yōu)化效果，但在2種擴展性測試中，各個測試規(guī)模下優(yōu)化后的HIP版本均比優(yōu)化前有明顯的性能提升。同時也可以發(fā)現，優(yōu)化后相比優(yōu)化前的加速比不及單結點時的情況。其主要原因是該算例PME部分計算量較多，使用CPU計算PME時該部分運行時間在總時間中占比始終過半，同時隨著結點數和算例規(guī)模的增加，PME部分通信量也迅速增加，對性能的制約更加嚴重，因此優(yōu)化效果逐漸降低。

圖7a和圖7b分別展示了弱擴展性和強擴展性測試中不同規(guī)模下優(yōu)化前后算例運行總時間和PME時間，可見優(yōu)化前后，隨著結點數的增長，PME部分時間占比都是增長趨勢，且優(yōu)化后，各個規(guī)模下PME時間占比都超過65%，說明PME以外的部分都得到了較好的優(yōu)化。

Figure 7 Percentage of PME time in scalability tests 圖7 擴展性測試PME時間占比

擴展性測試的結果表明，盡管本文并未針對GROMACS的通信以及計算流程等方面進行優(yōu)化，但通過對GPU核函數的優(yōu)化，在多結點計算時相對優(yōu)化前仍取得了可觀的性能提升。另一方面，該組測試也體現了PME部分對擴展性的限制，為后續(xù)針對多結點、大規(guī)模計算的優(yōu)化指明了方向。

4 結束語

本文通過將廣泛使用的分子動力學模擬軟件GROMACS 2020系列轉碼為HIP代碼，實現了其功能在ROCm上的完整移植。進一步以一個復雜的離子液體模擬算例為目標，對其主要GPU核函數進行了逐一優(yōu)化，實現了相對初始轉碼版本約2.8倍的加速比。優(yōu)化后的版本在MI50 GPU上的性能高于GROMACS原版OpenCL代碼60.5%，相對運行在EPYC 7502 處理器的原版CPU代碼的加速比約為2.7。此外，本文還在另外2個典型算例Lysozyme和waterbox上進行了測試。結果顯示，優(yōu)化后的HIP版本相對原版CPU和OpenCL版本均有一定的性能加速。通過對所用離子液體算例的擴展，分別測試了優(yōu)化前后HIP版本代碼的強弱擴展性，一方面說明針對單結點的優(yōu)化在多結點運行時仍然體現了良好的加速效果，另一方面也發(fā)現了PME部分對擴展性提升的制約，為后續(xù)針對大規(guī)模體系的優(yōu)化提供了方向。本文的工作豐富了GROMACS可使用的軟硬件平臺，所采用的移植方法、優(yōu)化思路及方法對于其他GPU應用的移植和優(yōu)化也有一定借鑒意義。未來我們將就更大規(guī)模的可擴展性優(yōu)化展開研究。