基于GPU加速的投影后變分殼模型計算

陸 曉,連占江,高早春

(中國原子能科學研究院 核物理研究所,北京 102413)

理論上,原子核多體波函數可通過求解薛定諤方程獲得。傳統殼模型就是在此基礎上建立的。然而,由于模型空間的限制,傳統殼模型計算僅限于質量數較輕或幻數附近的原子核。對于變形重核,需更大的價核子空間,相應的殼模型組態空間維度數會變得極其巨大,因而無法實現精確的殼模型計算。這是核多體問題中一直存在的難題,且在可見的未來難以被徹底解決。為回避巨大組態空間這一困難,人們發展了各種殼模型近似方法[1-13],以期進一步擴展殼模型應用核區范圍。

投影后變分(VAP)方法就是一類較重要的殼模型近似方法[8-13]。它將原子核試探波函數進行變分,從而得到盡可能好的殼模型近似結果。通常情況下,VAP波函數可由投影Hartree-Fock-Bogoliubov(HFB)真空態或投影Hartree-Fock(HF)斯萊特行列式(SD)構成。其中采用HFB真空態完全投影得到具有好量子數(中子數、質子數、角動量和宇稱)的波函數需要五重積分,計算非常耗時。而文獻[8]的研究表明,在這幾種投影中,角動量投影是最重要的。于是,在實際應用中,通常采用HF SD只進行角動量與宇稱投影構建初始試探波函數。顯然,在相同數目的投影態下,用投影SD作為VAP波函數的近似不如投影HFB真空態得到的結果好。但前者的近似性可通過添加更多的投影SD來改善。在某種意義上,采用投影SD進行投影后變分計算可減少計算消耗。因此,本文采用的VAP波函數全部基于投影SD構建。

然而,隨著VAP計算往中重核區推進,計算所需模型空間不斷擴大,價核子數目不斷增加,為保證計算的精度,就需更多的投影SD參與構建試探VAP波函數,這無疑會造成VAP計算負擔的加重,甚至導致VAP迭代收斂非常緩慢。隨著半導體和芯片技術的迅猛發展,可提供應用的高性能計算平臺不斷增加。特別是近年來計算性能飛速提升的圖形處理器(graphics processing unit, GPU),其在人工智能、大數據等領域應用廣泛,已逐漸取代早期用于并行計算的CPU處理器,成為現代高性能計算的首選。因此,在原子核物理領域,GPU卓越的計算性能也有望進一步提升相關物理模型的計算效率。作為一個具體實例,本文擬將VAP殼模型計算程序從傳統的CPU平臺移植到GPU平臺,以期進一步提升VAP的計算速度。

本文將從VAP的原理出發,介紹其在計算當中遇到的數值問題,基于OpenACC并行化模型實現VAP在GPU平臺上的計算,并通過實例展示采用GPU相對于CPU在計算效率上的提升。

1 VAP的數值問題

首先,采用投影SD構建的試探VAP波函數可寫為如下形式[12]:

(1)

(2)

若要使這m個態的波函數盡可能接近殼模型的精確波函數,需對|ΨJπMα(K)〉進行最優化。從式(1)可看出,|ΨJπMα(K)〉是由|Φi〉決定的。于是,最優化|ΨJπMα(K)〉實際上是通過對所有的|Φi〉變分實現。欲對|Φi〉進行變分,首先對其進行參數化。此過程可通過如下Thouless定理[14]實現:

(3)

DVAP=N(MN-N)+Z(MZ-Z)

(4)

在VAP計算中,將變分參數取為復數的形式。于是,采用n個SD構建的VAP波函數,其獨立變分參數就有2nDVAP個?？梢?隨著模型空間的增大、價核子數目的增多以及SD數目的增加,將會導致獨立變分參數迅速增加。

(5)

在最小化能量之和的過程中,為保證VAP計算的穩定性,在迭代過程中計算能量的梯度與Hessian矩陣[9]。即計算了如下4類矩陣元:

(6)

(7)

此外,在VAP計算中,哈密頓量通常取為如下形式:

(8)

其中:tij和vijkl分別為單體和兩體相互作用矩陣元;c+和c分別為粒子產生和湮滅算符;i、j、k、l分別表示不同的單粒子軌道。由式(8)可看出,在兩體相互作用下,單個投影矩陣元在特定積分格點處的分量通常包含一個對單粒子軌道的4重求和。

基于以上討論可看到,在VAP計算中,投影矩陣元的計算占據了大部分計算時間。并隨著模型空間的增大,價核子數的增多,計算時間會迅速增加。為更直觀體現這種計算耗時,以sd模型空間的24Mg與pf模型空間的56Ni的基態為例,采用1個SD,即n=1,對24Mg基態到56Ni基態投影矩陣元計算量的增加做一定量估計。首先,對24Mg的基態,每個方向上只需選取8個格點就足以得到好的積分結果。而計算56Ni的基態,每個方向上積分格點的數目通常需增加至18個。因此,從24Mg到56Ni,總的積分格點數目增加約了10倍。因為做簡單估算,將單個投影矩陣元在特定積分格點處每重求和的求和次數取為原子核所在模型空間的中子軌道數的一半。此時,24Mg所處的sd殼包含12條中子軌道,對應的總求和次數為1 296,而56Ni所處的pf模型空間則包含20條中子軌道,其總求和次數為10 000,接近24Mg求和次數的8倍。最后24Mg共有64個獨立的變分參數,對應有64個梯度分量和2 080個獨立的Hessian矩陣元,56Ni則共有192個獨立的變分參數,對應有192個梯度分量和18 528個獨立的Hessian矩陣元,分別是24Mg數目的3倍和8.9倍。因此,從24Mg的基態到56Ni的基態,總計算量增加了近780倍,且這一數字隨著模型空間的增大和價核子數目的增加會進一步增大。所以,如果采用傳統的串行VAP計算,在時間上是不可接受的。

在VAP中,式(6)所示的投影矩陣元之間是相互獨立的。因此,為節省計算時間,實際的VAP計算一般需借助并行技術實現不同獨立格點的同時計算。在之前的工作中,為提高中重核區的VAP計算效率,主要通過OpenMP等并行方法將不同積分格點處的轉動矩陣元計算任務分配給不同的CPU核心執行,從而實現對不同積分格點處轉動矩陣元的同時計算。在這種模式下,同一積分格點處的不同轉動矩陣元仍以串行方式逐個計算。從上面的討論可知,同一積分格點處存在大量相互獨立的轉動矩陣元。顯然,這些獨立的轉動矩陣元也可采用并行方式同時計算。因此,一種更有效的加速方式是將所有獨立的轉動矩陣元同時并行化。當然,這也意味著需更多的計算核心。在這方面,內嵌數千個CUDA內核的高性能GPU無疑較一般只有幾十或上百個核心的CPU計算服務器具更大優勢。因此,為進一步提升VAP的計算效率,有必要將當前的程序代碼移植到GPU平臺上進行加速計算。

2 基于OpenACC的GPU加速VAP程序改造

GPU是一種多核協處理器,通過PCIe總線與CPU相連。在并行計算中,程序總是從CPU開始執行,由CPU負責執行串行代碼并為并行計算做好準備,包括數據初始化、分配GPU內存、將數據復制到GPU等。當程序運行到計算密集區域時,CPU以特殊形式調用GPU,由GPU完成并行計算并將結果返回到CPU繼續向下執行。本文選用的GPU加速器為英偉達公司生產的Tesla V100。1塊Tesla V100 GPU共有5 120個CUDA內核和640個Tensor內核,不僅為Tesla V100提供了多核并行計算的結構基礎,也使其具備更高的浮點運算能力。目前,Tesla V100單精度浮點運算每秒峰值速度可達到15.7TFLOPS。

OpenACC是一基于指令的高級編程模型。編程人員只要在串行的C/C++或Fortran程序中插入合適的指令,支持OpenACC的編譯器就會根據指令的內容產生可執行的低級語言代碼,將指令指定范圍內的計算密集的代碼分派給GPU等加速器并行完成。對于不支持OpenACC的編譯器,指令會被當成注釋自動忽略。這種機制不僅有效地減少了對原程序的改動,也有利于程序的跨平臺運行。

基于OpenACC并行化模型可對VAP程序進行改造,將其移植到GPU平臺進行加速。由于VAP計算中的投影矩陣元之間是相互獨立的,因此,在將程序移植到GPU平臺的過程中,基本的思路就是利用OpenACC并行方法將同一積分格點處所有獨立的轉動矩陣元算出,然后對不同積分格點處的轉動矩陣元求和得到相應的投影矩陣元。

經GPU移植后的程序執行流程如圖1所示,圖中藍色程序塊表示串行區域,由CPU執行,橙色程序塊表示并行區域,由GPU執行。在程序執行過程中,首先由CPU完成初始化和預處理,再將初始基矢波函數和兩體相互作用信息從CPU內存拷貝到GPU內存開啟投影矩陣元的并行計算。待并行計算結束后,再將計算結果從GPU內存拷貝到CPU內存,在CPU上求解波函數滿足的HW方程,從而得到能量及其相應的梯度和Hessian矩陣,利用信賴域算法得到下一次迭代的初始波函數。重復以上過程直至能量達到極小。

圖1 OpenACC版本VAP程序流程圖

3 GPU加速效果檢驗

計算中選用的計算平臺為DELL PowerEdge T640服務器,每臺服務器同時搭載2顆Intel(R) Xeon(R) Gold 6148 CPU(20核|2.40 GHz) 和1顆英偉達公司生產的Tesla V100 GPU。

為檢驗經GPU移植后的VAP程序的正確性,從同一SD出發,分別用原來經過檢驗的OpenMP并行程序和新的OpenACC并行程序在sd殼對24MgJπ=0+、2+、4+、6+的低自旋態開展了完整的VAP計算。計算中所用的哈密頓量為USDB有效相互作用[16]。圖2給出了兩種情況下VAP能量隨迭代步數的變化。對于計算的所有自旋態,OpenACC并行程序每步得到的VAP能量與OpenMP并行程序得到的結果在計算機誤差范圍內完全一致。變分達到收斂所需要的迭代步數也完全相同。這證明經GPU移植后得到的OpenACC并行程序的正確性。

為進一步對比原來OpenMP并行程序和新的OpenACC并行程序的加速效果,分別計算VAP程序在CPU串行(單核)、CPU并行(40 核)和GPU并行3種條件下的單步運算時間。計算中只用1個SD,并固定迭代步數為10步。單步運算時間定義為每次VAP迭代中計算式(6)的所有投影矩陣元平均需要花費的時間。并行單步運算時間與串行單步運算時間的比值稱為加速倍數。

圖3給出了OpenMP并行程序和OpenACC并行程序在計算部分sd殼原子核基態時的加速倍數。從圖3可看出,經并行后,VAP程序的運算速度相比串行模式下有了較大的提升,且GPU并行的加速效果顯著優于CPU并行。對于28Si,GPU并行的加速倍數達到了接近50倍,而CPU并行的加速倍數只有不到11倍。此外,無論是采用CPU并行還是GPU并行,其加速倍數均隨價粒子數的增多而增大。不同的是,GPU并行的加速倍數的增長速度明顯高于CPU并行。從23Na到28Si,GPU并行的加速倍數增大了1.2倍,而CPU并行的加速倍數只增長了53%左右。這主要是因為sd殼模型空間較小,并不足以讓GPU核心得到充分利用。這一點也直接體現在VAP程序執行過程中的GPU核心占用率上。如在計算23Na時,GPU核心的占用率約60%,這意味著僅60%的GPU核心得到了利用。而在計算28Si時,這一數字達90%。因此可預見,隨著模型空間的增大和價粒子數目的增加,GPU并行的加速效果仍有望得到進一步提升,直至核心占用率達到100%。

圖3 CPU并行與GPU并行加速倍數隨不同原子核的變化

為進一步測試GPU并行的加速效果和驗證經GPU加速后的VAP程序在中重核區的適用性,在相同的平臺下對64Ge的基態開展了完整的VAP計算。采用fpg9/2模型空間,fpg9/2模型空間包含30條中子軌道和30條質子軌道。在此模型空間中,64Ge的組態空間維度達1.7×1014,較組態相互作用殼模型目前能解決的范圍高出3個數量級。在本次計算中仍然采用1個SD。在此情況下,GPU核心占用率達到了100%,加速倍數更是高達192倍。與CPU并行(40核)相比,單步運算時間從原來的75 min降到了9 min。

最后,作為一大規模殼模型計算應用實例,利用OpenACC并行程序對稀土區重核178Hf的基帶開展了計算。178Hf所處的jj56pn模型空間包含32條質子軌道和44條中子軌道。在此模型空間中,178Hf的組態空間維度達到了1.9×1018,較傳統的嚴格對角化方法的處理極限高7個數量級。處理如此大維度的計算問題,原來的OpenMP程序迭代1次需花費超過1 d的時間,而經過GPU加速的OpenACC并行程序迭代1次只需1.5 h左右。

圖4給出了采用一個SD進行VAP計算得到的178Hf基帶能譜。計算采用的相互作用為jj56pnb有效相互作用[17],實驗數據取自文獻[18]。從圖4可看出,通過VAP計算得到的178Hf基帶能譜與實驗值符合得很好,對應能級之間的能量差最大不超過400 keV。

圖4 178Hf基帶能譜

4 總結與展望

VAP方法作為一種重要的殼模型近似方法,有望在較大模型空間中實現對中重核區原子核的計算。但在向中重核區推進的過程中,隨著模型空間的增大以及價核子數的增多,VAP需計算的投影矩陣元的數目會迅速增加,極其耗費計算時間。由于投影矩陣元之間彼此是相互獨立的,因此,VAP計算可通過并行化處理以提高計算效率。隨著近年來高性能GPU計算平臺的迅速發展,相較于原來在CPU平臺上采用OpenMP并行化程序進行VAP計算,在高性能GPU平臺上采用OpenACC并行化模型進行計算有望進一步提升計算效率。通過將VAP程序采用OpenACC并行化模型改造,在角動量投影的每個積分格點上,實現了數目龐大的各獨立轉動矩陣元的GPU并行化計算。經驗證,GPU加速后的VAP程序與原來的OpenMP并行化程序相比,計算效率進一步提升了數倍。且隨著往重核區的推進,提升倍數會進一步增加。采用改造后的程序,首次實現了稀土區重核178Hf的基帶能譜的計算。目前,VAP程序只是在1塊GPU內進行了加速,其計算效率即得到顯著提升?？深A見,若采用多塊GPU同時對VAP程序進行加速,將會進一步提升VAP的計算效率,使得VAP計算更易往重核區推進。這一工作將會在未來展開。

GPU加速的VAP計算程序首次具備了在GPU服務器單機開展重核區計算的能力,不僅節省了大量的能耗,同時為系統研究中重質量核區復雜原子核結構問題提供了極大的便利。過去,重質量原子核的大規模殼模型計算即使在超級計算機集群上也難以開展起來,對重質量原子核結構的系統研究普遍缺乏微觀的殼模型基礎。目前對于許多復雜核結構現象的理解還處于比較模糊的認識階段,如γ振動帶與β振動帶的同時描述,復雜的奇A核與奇奇核八極轉動帶,超形變帶的帶間躍遷等。對這些現象的深入理解必須借助于可靠的微觀殼模型方法。本文建立的VAP殼模型計算將有望在這些問題上發揮其獨特作用。