趙文龍，王武

1.中國科學院計算機網(wǎng)絡信息中心，北京 100083

2.中國科學院大學，北京 100049

引言

N 體問題描述了多個粒子在經(jīng)典力學下的運動規(guī)律，是廣泛運用于天體物理學和分子動力學等多個重要領域中的基本問題之一。為了能夠更加精準地預測宇宙中物質(zhì)的分布以及預想模型的正確性，往往需要進行大規(guī)模的模擬，這對于傳統(tǒng)計算機集群來說有極大的壓力。而隨著硬件與算法的快速發(fā)展，用來求解N 體問題的軟件也在不斷地更新迭代，如GreeM[1]、CUBE[2]和PHoToNs[3]。

目前模擬N 體問題的主要數(shù)值方法有粒子-粒子法（Particle-Particle,P2P），粒子-網(wǎng)格法[4]（Particle-Mesh Method,PM），Barnes-Hut 樹算法[5]（Barnes-Hut Tree Method,BH-Tree）以及快速多極子算法[6]（Fast Multipole Method,FMM）。P2P 方法的時間復雜度為，計算效率較低，但是精度最高；PM 方法將空間網(wǎng)格離散，通過快速傅里葉變換（Fast Fourier Transform,FFT）求解泊松方程得到作用力，但近距離時誤差較大；而樹方法，包括BH-Tree 算法與FMM 算法，可以控制計算精度。為了兼顧精度與效率，P3M[7]、TreePM[8]與FMM-PM[9]這樣的混合算法成為了主流。

得益于目前高性能計算的快速發(fā)展，尤其以CPU+GPU 為主流的并行異構平臺的強大的計算能力，N 體問題成為了加速卡異構平臺的熱點應用之一。

國際上已有不少針對GPU 這種SIMT（Single Instruction Multiple Thread）并行機制的加速卡的N體模擬異構算法與實現(xiàn)。Nyland[10]針對P2P 算法的優(yōu)化提出了“i-parallel”方法以及“walk”的概念，在GPU 上加速了50 倍。而Yokota[11]在Nyland的基礎上，實現(xiàn)了對107個粒子的150 倍加速。由Hamada 提出的“Chamomile Scheme”算法被完全移植到了GPU 上，最終模擬131,072 個粒子性能達到了256GFlops[12]，之后通過對樹算法進行一系列優(yōu)化，用576 個GPU 模擬32 億個粒子，性能達到了190.5TFlops[13-15]。2011 年，波士頓大學針對FMM算法在天文問題上的應用進行GPU 加速[16]，相較于CPU 端加速了70 倍。2017 年，Potter 將其開發(fā)的PKDGRAV3[17]中的FMM 算法成功移植到GPU 上，并用4,000 個GPU 節(jié)點實現(xiàn)了2 萬億粒子的天文模擬。在國內(nèi)，針對國家天文臺研發(fā)的PHoToNs-2 有相關的多GPU 的移植與優(yōu)化[18-19]。上海交通大學以及廈門大學對CUBE 的PP 部分進行GPU 移植，單步運行速度提升了5.16 倍，并在64 個GPU 上有著83.8%的弱擴展效率。

Gadget-2[20]是天體物理領域最有影響力的并行模擬軟件之一，它由Springel 開發(fā)，采用TreePM 算法，是國際上第一個可計算粒子規(guī)模達到百億量級的開源宇宙模擬軟件，在天文領域獲得廣泛認可，相關論文被引用超過6,500 次。針對不同的計算平臺，Gadget 系列已經(jīng)衍生出了眾多分支版本來適應不同的并行架構和編程模式，如OpenAcc-Gadget3[21]、MP-Gadget[22]、AX-Gadget[23]。在2021 年推出的Gadget-4[24]采用了FMM-PM 算法。

本文針對國產(chǎn)高性能異構加速卡平臺，基于HIP 完成了Gadget-2 的移植，并對其在核函數(shù)的訪存效率進行了優(yōu)化，最終整體性能加速13.27 倍，短程力計算加速35.67 倍，并行效率達到57.29%。第1 節(jié)主要介紹TreePM 以及相關算法的基本原理；第2 節(jié)介紹加速卡平臺的計算環(huán)境，以及最耗時的短程力計算在加速卡上的實現(xiàn)方法；第3 節(jié)介紹短程力計算在加速卡上的性能優(yōu)化方法；第4 節(jié)驗證移植優(yōu)化版本的結果正確性，并給出性能測試結果。

1 TreePM 算法原理

1.1 BH-Tree 算法

BH-Tree 算法是由Barnes 和Hut 提出的，其主要思想是將模擬空間看作是一個立方體，每一個維度遞歸地劃分子空間，將空間中的粒子劃分到不同的子節(jié)點，直到每個子空間只包含了一個粒子，以此構造一個八叉樹，并將模擬的整個空間看作是根節(jié)點。

構造樹之后，從葉子節(jié)點的父節(jié)點開始，由下往上依次計算每個節(jié)點的質(zhì)點與質(zhì)量，直到根節(jié)點。若節(jié)點只包含一個粒子，則該節(jié)點的質(zhì)量與質(zhì)心由該粒子的質(zhì)量與位置決定，否則由子節(jié)點決定。在計算作用力的時候，自根節(jié)點向下遍歷，粒子之間的相互作用轉(zhuǎn)換為包含多個粒子的樹節(jié)點與粒子之間的相互作用，這使得計算復雜度降低為，是否能夠?qū)鄠€粒子的節(jié)點看作一個整體來對目標粒子進行作用，是通過以下準則來判定：

圖1 BH-Tree 算法中八叉樹的劃分及粒子受力Fig.1 Partitioning of octrees and the forces acting on the particles in the BH-Tree algorithm

1.2 粒子-網(wǎng)格算法

粒子網(wǎng)格算法是將空間離散網(wǎng)格化，先計算引力勢然后再通過引力勢計算其所受作用力。粒子的引力勢滿足泊松方程[25]：

1.3 TreePM 算法

TreePM[26]算法由Xu 等人提出，其主要思想是由于PM 的運算效率比樹方法快得多，但同時在小尺度范圍內(nèi)計算精度較低，而樹方法可靈活控制精度但計算大尺度范圍相互作用的效率較低，模擬大規(guī)模N 體問題依然存在性能瓶頸，因此綜合權衡兩個算法的優(yōu)點和不足之處，TreePM 算法將引力勢在譜空間中劃分為兩部分：長程力與短程力，樹方法計算短程力，PM 方法計算長程力[27]：

實現(xiàn)時并不直接計算公式中的超越函數(shù)，而是通過樣點查表和插值多項式快速計算

TreePM 方法與P3M 方法相比誤差近似甚至更小，并且在樹方法中只需要遍歷目標粒子周圍的一個小空間，兼顧了數(shù)值計算效率，使得該算法在多個領域中獲得了廣泛的應用，帶來了巨大的性能提升。

但即便如此，TreePM 的性能瓶頸隨著模擬尺度與粒子規(guī)模的增大依然明顯。經(jīng)過測試，PM 部分通過FFT 計算，計算時間在整個算法的運行過程中占用比例極小，而樹方法中計算短程作用力的部分是計算密集型的任務，密度分布較集中時短程力計算時間平均占比90%左右，因此通過高性能異構平臺對其進行硬件加速，利用多線程計算來提高這部分的計算速度，能夠有效地提升整個算法的求解效率。

2 TreePM 在加速卡上的實現(xiàn)

2.1 ROCm 計算平臺與HIP 編程模型

“中科先導1 號”是基于國產(chǎn)處理器和加速卡的高性能異構平臺，其基本架構如圖2 所示，其中每個節(jié)點有一個32 核的CPU 處理器以及4 張類GPU 加速卡，分別為Host 端與Device 端。Host 端的內(nèi)存為128GB，每張加速卡有64 個流多處理器和16GB 顯存，訪存帶寬1TB/s，與CPU 之間使用32 個PCIE 3.0 互聯(lián)。加速卡與CPU 之間的訪存帶寬是16GB/s，節(jié)點間的網(wǎng)絡帶寬是25GB/s。CPU 支持x86 指令集以及常用的并行編程模型如OpenMP、MPI 等，加速卡支持ROCm 計算環(huán)境和HIP 編程模型。

圖2 “中科先導1 號”架構Fig.2 Architecture of Xiandao-1,CAS

ROCm[28-29]開源計算平臺基于新一代的Radeon硬件，支持最新的數(shù)學函數(shù)庫和匯編語言環(huán)境，以及多加速卡的計算，為開發(fā)大規(guī)模、高性能計算軟件提供了便利，并大大降低了功耗，同時該平臺還提供了機器學習框架等多種開源技術，并在不斷地更新，有著優(yōu)秀的可擴展性，其中用于可移植性的異構計算接口HIP 是一種C++運行時API 和內(nèi)核語言，它主要是為了簡化CUDA 應用向可移植C++代碼的轉(zhuǎn)換，接口可以跨主機/內(nèi)核邊界使用模板和類，以及其他C++功能集。移植工具HIPify 可以自動完成大部分CUDA 到HIP 的轉(zhuǎn)換工作，代碼在AMD 的DCU 與NVIDIA 的GPU 上通過各自的編譯器進行編譯后便可以正常運行，同時沒有性能上的損失。

2.2 樹方法的短程力計算在加速卡上的實現(xiàn)

TreePM 方法中，短程力部分采用樹方法計算。加速度的計算是整個N 體模擬計算流程最主要的開銷，占比可以達到90%以上，而這一部分中計算短程力的部分占了極大比例。以表1 為例，粒子數(shù)為1283，采用4 個進程，迭代200 個時間步。加速度的計算耗時905.1 秒，占總時間的93.6%，而其中計算短程力的部分可以占總耗時的89%，因此對這一部分進行加速，可以大幅度地提升整個流程的效率。

表1 Gadget-2 的各部分耗時Table 1 Time consumption for each part of Gadget-2

在Gadget-2 中，樹方法的實現(xiàn)是一個“邊走邊算”的過程，即獲得目標粒子后，從根節(jié)點開始，對目標粒子附近的空間節(jié)點進行遍歷，獲得相關的作用節(jié)點，然后計算兩者的作用力，因此本文樹方法在加速卡上的實現(xiàn)中，將對樹遍歷的過程與粒子作用力計算同時轉(zhuǎn)移到設備端進行操作，相應的我們需要將粒子信息與樹的相關信息共同傳入設備端，即如圖3 所示。

圖3 Tree 方法在加速卡的實現(xiàn)Fig.3 Implementation of BH-Tree method on accelerators

短程力計算中，首先在CPU 端將樹節(jié)點的相關信息傳入設備端，再通過函數(shù)CountParticlesInCurrTime()遍歷粒子列表，得到當前時間步的目標粒子數(shù)量以用于分配傳送粒子信息所需要的內(nèi)存。

開始對粒子進行引力計算。對短程力的計算包含兩個部分：local 和export，分別表示當前進程以及其他進程的粒子對目標粒子進行的作用力。在計算local list 時，遍歷樹的過程中會判定并標記當前粒子是否會對其他進程的粒子進行作用，當計算完本地進程粒子之間的作用力之后，根據(jù)標記從其他進程獲取需要進行相互作用的粒子列表，開始計算export list。實現(xiàn)時在函數(shù)force_treeevaluate_shortrange_hip()中調(diào)用核函數(shù)進行引力計算，并傳入?yún)?shù)0 或者1 來表示計算local list 或是export list，從而在核函數(shù)開始前進行對應的數(shù)據(jù)預處理。最終計算完成之后對結果進行匯總，傳回CPU 端后，更新粒子加速度以進行下一次迭代。

在原本的Gadget-2 中，是將當前時間步中參與引力計算的粒子劃分為多次任務后進行計算，為了減少核函數(shù)的調(diào)用次數(shù)以及數(shù)據(jù)傳輸開銷，在計算第一次任務時即可將當前進程中所有粒子的本地加速度在加速卡上計算完成，之后的多次任務只進行export list 中的粒子的引力計算。

核函數(shù)中線程根據(jù)索引取得對應的粒子信息，然后通過遍歷樹計算粒子與節(jié)點之間的作用力。遍歷從根節(jié)點開始，其流程與原Gadget-2 基本相同，都是依靠NextNode 數(shù)組來獲得下一個需要遍歷的節(jié)點，然后進行判定處理，計算加速度并更新。由于我們需要在遍歷樹的時候標記粒子是否會對其他進程中的粒子進行作用，此時粒子在線程中的索引與其在CPU 端Particle 數(shù)組中的索引是不同的，因此還需要將粒子在CPU 端實際的索引傳到設備端。Gadget-2 在構造樹時高效利用了節(jié)點之間的索引關系，相關內(nèi)存的消耗大大降低，并且提高了遍歷樹的效率，同時也有效地減少了需要傳入設備端的數(shù)據(jù)量。經(jīng)過測試，在模擬規(guī)模達到1283粒子的整個過程中，數(shù)據(jù)傳輸所占用的時間比例僅是2.5%。

3 核函數(shù)在加速卡上的優(yōu)化

對于加速卡來說，核函數(shù)的內(nèi)存請求通常是在DRAM 設備和片上內(nèi)存間以內(nèi)存事務來實現(xiàn)的。連續(xù)對齊、合并的內(nèi)存訪問對于設備內(nèi)存來說是最佳訪問模式，否則會因為內(nèi)存事務中有大部分傳輸?shù)臄?shù)據(jù)不會被使用，而造成帶寬的浪費，或者是事務請求的數(shù)量過于龐大，增加了訪存壓力，降低了最終的訪存效率。通常來說，有兩種數(shù)據(jù)組織的方式，分別是Struct of Array (SOA)與Array of Struct(AOS)。相比于AOS，SOA 通常能更加有效地減少訪存事務的數(shù)量，提高訪存效率。

如圖4，采用AOS 模式在加速卡上讀取坐標數(shù)據(jù)x 字段，由于隱式加載以及數(shù)據(jù)的交叉存儲訪問，每次讀取事務中只有1/3 的數(shù)據(jù)是我們實際需要的，同時會由于加載了不需要的數(shù)據(jù)而浪費緩存空間。如果線程0 到線程5 需要讀取x 軸坐標，按照圖示讀取方式需要三次事務請求，而在SOA 模式下讀取數(shù)據(jù)只需要一次事務請求，更好地利用了加速卡的內(nèi)存帶寬，以達到更高效的訪存效果。

圖4 SOA 與AOS 內(nèi)存布局Fig.4 The memory layout of SOA and AOS

如算法2 中的偽代碼所示，在核函數(shù)中，在遍歷樹的過程中會反復用到許多常數(shù)來判斷節(jié)點之間的關系，同時在更新加速度因子fac 時，會涉及到對表的查詢，該表是對修正系數(shù)的插值。為了提高線程的訪存效率，在不影響并行效率的情況下這些常數(shù)可以先預取保存在寄存器中，而對于查詢表可以將其保存在共享內(nèi)存。

優(yōu)化過程中也可以采用頁鎖定內(nèi)存以及異步多流的方式來提高數(shù)據(jù)在CPU 端與設備端之間的傳輸效率，但是如2.2 節(jié)所述，數(shù)據(jù)傳輸占據(jù)的整體運行時間比例僅為2.5%，因此該類方法并不能大幅提升整體性能。

4 數(shù)值結果與正確性

4.1 正確性驗證

下面將Gadget-2 的HIP 移植優(yōu)化版本記作G2H。實驗平臺為“中科先導1 號”，軟硬件環(huán)境詳見2.1 節(jié)。宇宙模擬的初始條件由基于二階拉格朗日擾動理論和Lambda Cold Dark Matter（LCDM）模型編寫的軟件2LPTic 生成。LCDM 模型是宇宙大爆炸理論的數(shù)學參數(shù)化描述，該模型假設宇宙中的大部分物質(zhì)由冷暗物質(zhì)組成，它們只存在引力相互作用。這是如今宇宙學中普遍認同的標準化模型，也是目前可以為觀察到的宇宙微波背景輻射、宇宙的大尺度結構等現(xiàn)象提供合理解釋的最簡單的模型，同時在該宇宙模型中有一個與暗能量相關的常數(shù)Lambda，它能夠解釋哈勃定律，即宇宙隨時間不斷加速膨脹的過程。

文中用功率譜來驗證G2H 的計算正確性。驗證方法是通過計算統(tǒng)計暗物質(zhì)密度分布，然后根據(jù)模擬尺度得到密度分布對應的功率譜。功率譜是隨機振幅的兩點相關函數(shù)的傅里葉變換，這個振幅會在引力的影響下隨著時間的推移而緩慢變化，因此我們可以通過比較程序模擬中相同紅移處的功率譜檢驗正確性。

實驗過程中模擬1283個粒子在尺寸為100h-1Mpc，以共動坐標系為參考的盒子中的運動狀態(tài)，并考慮周期邊界條件，初始紅移z=49，其他參數(shù)包括：同時采用PM 方法，網(wǎng)格參數(shù)設置為PMGRID=128。實驗過程中選取紅移z=0、1、2、3、5 來進行功率譜的比對。

Gadget-2 和G2H 在紅移z=0 時粒子分布的柱密度圖如圖5 所示，圖6 的功率譜中線條和符號分別表示G2H 和Gadget-2 的結果，兩者的殘差幾近為0，保證了G2H 的正確性。

圖5 Gadget-2（左）與G2H（右）的N 體模擬在紅移z=0 時的柱密度圖Fig.5 The density map for the N-body simulation of Gadget-2(left panel) and G2H (right panel) as red shift is 0

圖6 Gadget-2 與G2H 的功率譜比較Fig.6 Comparison of power spectrum of Gadget-2 and G2H

4.2 性能測試

下面分別測試加速比和并行效率。由于計算過程中每一步粒子的分布不同，因此每個時間步的運行時間有波動，我們進行10 步迭代然后取平均值作為單步運行時間進行比較。粒子規(guī)模分別取1283、2563、5123，分別使用4、8、16 個進程(記作Np)與相同數(shù)量的加速卡對Gadget-2、G2H 進行性能測試，平均單步的總運行時間和短程力運行時間見表2。

由表2 可以得到，粒子規(guī)模越大，程序在加速卡上的加速效果越明顯，因為短程力的計算時間占比一般會和粒子規(guī)模一同增加。相比于Gadget-2，G2H 程序的整體加速比可以從11.47 倍提高到13.27倍，短程力部分計算加速比可以從24.25 倍提高到35.67 倍。

表2 Gadget-2、G2H 的單步平均耗時Table 2 Average step runtime of Gadget-2 and G2H

在以往PHoToNs 的移植優(yōu)化工作中，短程力計算效率可以提升上百倍，而Gadget-2 中熱點函數(shù)在加速卡上移植優(yōu)化后性能只能提升30 倍左右，這與兩者的計算流程有關。在PHoToNs 中提前進行樹的遍歷，將相互作用的粒子放入包中，而加速卡提取對應的包進行計算，沒有過多的其他步驟，但在Gadget-2 中由于需要通過遍歷樹不斷判定來進行引力對的計算，除計算外的冗余步驟以及邏輯判斷較多，這些都嚴重影響了G2H 中核函數(shù)的并行效率，導致兩者的優(yōu)化效果有著明顯的差異。

在并行效率測試中，粒子規(guī)模為5123，分別使用8、16、32、64、128 個進程與相同數(shù)量的加速卡迭代10 步后取平均值作為單步耗時。圖7 顯示了G2H 在“中科先導一號”上的強可擴展性，表3 是測試中對應的具體耗時?？梢钥闯?，在加速卡上不僅有短程力的計算，還有樹的遍歷，這對短程力在加速卡上的并行效率有著較大的影響，而在主機端實現(xiàn)了非本地樹節(jié)點之間的通信以及PM 的通信，這些因素也成為了程序主要的性能瓶頸。最終G2H 的整體并行效率是57.29%，其中短程力并行效率是53.87%。

表3 G2H 強可擴展性測試的單步平均耗時Table 3 Average step runtime of G2H in strong scalability test

圖7 G2H 的強可擴展性測試Fig.7 The strong scalability test of G2H

5 結論與展望

本文將基于TreePM 算法的宇宙學N 體模擬軟件Gadget-2 中耗時占比90%以上的短程力計算模塊通過HIP 移植到加速卡異構平臺“中科先導一號”。通過數(shù)據(jù)內(nèi)存重構，并充分利用寄存器和共享內(nèi)存，減少了加速卡對內(nèi)存存取開銷，提高了設備端訪存效率。通過功率譜驗證了移植優(yōu)化的正確性，并進行相關的性能測試，程序整體與短程力部分的計算速度最高分別提升了13.27 倍與35.67 倍，并行效率分別達到57.29%與53.87%。由于本文的方案是在加速卡上進行樹的遍歷，導致較多的邏輯判斷，這對加速卡性能的充分挖掘是不利的，因此考慮將樹的遍歷盡可能地在CPU 上進行，在加速卡端進行主要的計算，同時這樣也能平衡主機與設備端之間的數(shù)據(jù)傳輸，這些將作為下一步的研究任務。本文的移植優(yōu)化方法可為類似的樹結構算法在高性能異構平臺上的實現(xiàn)與大規(guī)模異構并行模擬提供借鑒。

利益沖突聲明

所有作者聲明不存在利益沖突關系。