第一性原理軟件北京原子技術(shù)模擬工具包的重構(gòu)進展*

馬英晉 張?zhí)?何連花? 金鐘

1) (中國科學(xué)院計算機網(wǎng)絡(luò)信息中心, 北京 100190)

2) (中國科學(xué)院計算科學(xué)應(yīng)用研究中心, 北京 100190)

3) (北京凝聚態(tài)物理國家研究中心, 中國科學(xué)院物理研究所, 北京 100190)

4) (中國科學(xué)院大學(xué)物理科學(xué)學(xué)院, 北京 100049)

基于第一性原理的理論方法的研究, 代表了材料計算、分子模擬等領(lǐng)域的科學(xué)高地, 相應(yīng)的第一性原理計算軟件直接關(guān)系到該領(lǐng)域相關(guān)理論、算法的積累.本文匯報了我們在重構(gòu)第一性原理計算模擬軟件—北京原子技術(shù)模擬工具包(BSTATE)的一些最新進展.重構(gòu)的核心思想是降低用戶使用門檻、擴展軟件適用范圍、增加軟件對于流行計算框架的支持.基于此思路, 在BSTATE原有Makefile編譯系統(tǒng)的基礎(chǔ)上添加了CMake編譯環(huán)境, 并支持各種數(shù)學(xué)函數(shù)庫的自動和交互式配置; 通過在原有內(nèi)置泛函基礎(chǔ)上增添Libxc泛函庫的支持, 使BSTATE支持的泛函數(shù)量有了數(shù)量級上的增長; 分析測試BSTATE在集群的并行特點, 并以更新數(shù)學(xué)庫接口(FFTW3、Cufftw)的形式提供對于流行異構(gòu)框架的初步支持.

1 引 言

電子結(jié)構(gòu)計算的理論依據(jù)是量子力學(xué)理論.量子力學(xué)建立了微觀粒子運動變化的基本規(guī)律, 將人類對自然的認識擴展到原子尺度的微觀世界.非相對論量子力學(xué)最流行的表述形式是Schr?dinger的波動力學(xué)形式, 它的核心是波函數(shù)的概念及其運動方程 Schr?dinger方程.但電子體系的Schr?dinger方程是高維奇異問題, 很難直接求解.于是人們又通過各種手段對Schr?dinger方程進行等價轉(zhuǎn)換或近似, 其中最基本的近似是絕熱近似或稱Born-Oppenheimer近似.它的基本思想是將整個問題分成兩部分來考慮, 考慮電子運動時原子核是處在它們的瞬時位置上, 而考慮核的運動時則不考慮電子在空間的具體分布.目前實際計算中用于求解Schr?dinger方程的各種方法大都是在這個近似基礎(chǔ)上進行的.根據(jù)是否使用經(jīng)驗參數(shù)和實驗數(shù)據(jù),求解Schr?dinger方程的近似方法主要分為3類:第一性原理方法、半經(jīng)驗方法和經(jīng)驗方法.第一性原理方法是指不使用任何經(jīng)驗參數(shù)和實驗數(shù)據(jù), 只利用基本物理常數(shù)和原子量計算Schr?dinger方程的方法, 求解后得到體系的一系列性質(zhì), 如晶體結(jié)構(gòu)、結(jié)合能和電子結(jié)構(gòu)等.

眾所周知, 化學(xué)和材料科學(xué)是現(xiàn)代科學(xué)技術(shù)和國民經(jīng)濟發(fā)展的重要基石, 基于第一性原理的分子、材料模擬是研究化學(xué)反應(yīng)和材料物性最重要的方法之一, 為目標功能導(dǎo)向的材料設(shè)計提供了可能.第一性原理計算方法在材料的性質(zhì)研究和預(yù)測方面發(fā)揮著越來越大的作用, 已經(jīng)成為繼實驗和理論分析之外的最為重要的研究方法和手段之一.尤其是對于新型的拓撲材料的預(yù)測和研究, 第一性原理方法發(fā)揮了無可替代的作用.

當前我國在材料計算、反應(yīng)模擬領(lǐng)域仍然缺乏核心的競爭力和影響力, 一個重要的原因就是成熟的第一性原理系列計算軟件的缺失, 仍處于利用國外 成 熟 商 業(yè) 軟 件 , 如 : GAUSSIAN[1], VASP[2](Vienna Ab-initio Simulation Package, 維也納從頭算模擬程序包), 進行研究工作的下游.這也導(dǎo)致相關(guān)基礎(chǔ)人才儲備不足、后勁缺乏, 影響到計算研究進一步的發(fā)展.雖然目前國內(nèi)的許多研究小組已經(jīng)在該領(lǐng)域做出了出色的工作, 也有越來越多的相關(guān)程序包發(fā)布 (如 BDF[3], ABACUS[4], BSTATE[5]等), 但是目前有關(guān)這方面的研究積累仍然不足, 尤其是國產(chǎn)的適用于大型計算集群、支持異構(gòu)框架的大規(guī)模并行軟件仍然缺失.這極大限制了可計算模擬的材料尺度或者分子大小, 進而制約了相關(guān)材料和化學(xué)科學(xué)的發(fā)展.

BSTATE (Beijing Simulation Tool for Atom Technology), 中文名為北京原子技術(shù)模擬工具包,是基于密度泛函理論和平面波贗勢(包括模守恒贗勢[6]和超軟贗勢[7,8])方法的第一性原理計算程序包, 主體程序在20世紀90年代由中國科學(xué)院物理所方忠研究員及其合作者共同開發(fā).BSTATE主要用于周期性晶體結(jié)構(gòu)下能量、電子密度、能帶結(jié)構(gòu)等的計算模擬, 進而得到物質(zhì)的電、磁、光等其他特性.除了可以很好地處理一般的基態(tài)電子結(jié)構(gòu)計算, 該程序包的長處在于能比較好的處理各種磁有序、軌道序、強自旋軌道耦合體系, 包含了獨創(chuàng)的處理強關(guān)聯(lián)體系的LDA+Gutzwiller[9?11]方法,大幅提高了計算效率[5].基于BSTATE程序包, 以物理所凝聚態(tài)研究組為主體的研究人員針對Bi2Te3, Bi2Se3和Sb2Te3等拓撲材料進行了第一性原理的計算, 其研究結(jié)果表明這些體系是首次發(fā)現(xiàn)的三維拓撲絕緣體, 并且在其二維磁性摻雜薄膜中能實現(xiàn)量子化的反常霍爾效應(yīng).這也為極低能量耗散的自旋電子器件設(shè)計指出了一個新的發(fā)展方向[12?15].

雖然BSTATE程序包已經(jīng)在拓撲絕緣體的計算中取得了很好的成績, 但是BSTATE的使用主要局限在少數(shù)課題組內(nèi), 有必要進行進一步的推廣.當前BSTATE局限性的主要在于用戶門檻較高; 泛函數(shù)量較為單一, 不好滿足特定材料的計算要求; 以及暫時沒有考慮流行的硬件框架.針對上述問題, 我們擬對中科院物理所的第一性原理材料計算軟件包BSTATE進行重構(gòu).以期提高軟件的易用性和適用范圍、降低用戶門檻、優(yōu)化其對流行計算框架的支持.

2 BSTATE程序的基本原理

BSTATE程序為大型的第一性原理程序包,有近30萬行的Fortran代碼.程序使用了FFTW(the Faster Fourier Transform in the West)函數(shù)庫、Intel MKL (Math Kernel Library)等多種數(shù)學(xué)函數(shù)和 MPI (Message Passing Interface)來提高計算效率.程序的核心部分是對交換關(guān)聯(lián)勢做了 LDA[16](Local Density Approximation)或 者GGA[17,18](Generalized Gradient Approximations)近似的KS (Kohn-Sham)方程的自洽求解[19].

2.1 Schr?dinger方程與DFT理論

波恩-奧本海默近似下, 體系中電子的定態(tài)可以由滿足非含時Schr?dinger方程的N電子波函數(shù) Ψ (r1,···,rN) 來描述, 即

其中, H為體系哈密頓算符, 包含動能T(拉普拉斯算符 ?)、勢能V、電子相互作用U 三部分.由于電子之間的相互作用會讓體系成為很難處理的多體問題, 無法精確求解.DFT可以將含U的多體問題轉(zhuǎn)化為不含U的單體問題, 從而成為解決此類問題的一個有效方法.密度泛函理論中, 關(guān)鍵的變量為粒子的密度, 即

(2)式的關(guān)系也可以反過來, 即給出電子的密度ρ(r), 原則上也可以計算出對應(yīng)的波函數(shù) Ψ .基于此, Hohenberg與Kohn(HK定理[20])指出存在電子密度 ρ (r) 的總能量泛函 E (ρ) , 其中 E (ρ) 可以寫成外勢 Vext的形式, 即

其中, Vext由 ρ (r) 唯一確定, J [ρ] 為經(jīng)典的庫倫排斥能.此時可以通過最小化 E (ρ) , 來獲得體系的能量、電子密度以及其他所有性質(zhì).相比于計算量隨粒子數(shù)指數(shù)次增長的波函數(shù)方法, 電子密度方法的計算量與電子數(shù)目無關(guān), 僅與在三維空間 (x, y, z)內(nèi)找到電子的概率有關(guān).由于泛函 G (ρ) 是未知的,HK定理并沒有給出計算 E (ρ) 的實際方法.

2.2 KS-DFT與泛函近似

KS理論下 G [ρ] 可以表示為 Ts[ρ]+Exc[ρ] 的形式[21,22], 其中 Ts[ρ] 是假想無相互作用體系的動能泛函, Exc[ρ] 是交換關(guān)聯(lián)能量.Exc[ρ] 的求解需要交換關(guān)聯(lián)勢(νxc(r)= δExc[ρ]/δρ(r))的近似形式, 其最簡單的近似就是局域密度近似, 即LDA.LDA近似下 ρ 用均勻電子氣模型來處理.LDA交換關(guān)聯(lián)能只與 ρ 有關(guān), 表示為

其中, εxc(ρ(r)) 是均勻原子氣單粒子的交換關(guān)聯(lián)能.在LDA中, 具有相反自旋的電子占據(jù)相同的KS空間軌道.

為了校正電子密度分布不均勻帶來的誤差, 將表示電子密度不均勻性的密度梯度(? ρ)作為半局域變量, 引入到 Exc, 此為含廣義密度梯度校正的近似即 GGA, 得到的泛函即為GGA泛函.GGA 交 換 關(guān) 聯(lián) 能是 ρ 和 ? ρ 的 泛 函 , 即考慮到電子自旋后GGA泛函需要由四個參數(shù)確定, 即也因此相比于LDA泛函, GGA泛函在計算結(jié)合能方面明顯提高了計算精度.

除 了 LDA和 GGA泛 函, meta-GGA和hybrid-GGA[23,24]以及雙雜化泛函都是重要的泛函形式.當前BSTATE暫時沒有用到后三種形式的泛函.泛函確定后, 體系的波函數(shù)需要通過求解Hφ=Eφ形式的KS方程來得到, 具體形式為

其中 νeff為泛函確定后的有效勢.

2.3 BSTATE程序框架

作為第一性原理計算程序包, BSTATE程序包的核心部分是對交換關(guān)聯(lián)勢做了LDA或者GGA近似的KS方程的自洽求解.其基本流程示意圖如圖1所示.

圖1 BSTATE 計算流程示意圖Fig.1.Flowchart of BSTATE package.

以上數(shù)值求解的過程中必須把連續(xù)量離散化,如系統(tǒng)實空間原胞劃分成足夠細的網(wǎng)格, k空間的布里淵區(qū)的k點離散化.在迭代求解KS方程的過程中需要初始化電荷密度 ρ 與KS軌道, 可以利用單原子的電子密度分布作為初始值.有HK定理保證, KS勢可以表示成電荷密度的泛函; 所以很容易得到離散網(wǎng)格上的KS勢, 進而可以求解固定KS勢場的單電子本征方程獲得KS軌道信息.一般情況, 求解出來的KS軌道與初始KS軌道是不相同的.軌道求解出來以后可以根據(jù)DFT波函數(shù)為單行列式的性質(zhì)求得電子密度信息 ρ .獲得新的電子密度后, 使用HK定理重新計算KS勢, 重新求解KS方程獲得KS軌道信息, 最后重新利用KS軌道和電子排布獲得電子密度信息.更新重復(fù)上面的過程, 直到所求出的新的電荷密度與輸入的電荷密度在判據(jù)范圍內(nèi)一致, 即可認為收斂.這時候體系的波函數(shù)信息包括KS軌道和電子占據(jù), 在此基礎(chǔ)上可以得到系統(tǒng)的總能、電子態(tài)密度、能帶圖、布洛赫波、費米面、光電導(dǎo)等一系列結(jié)果.

3 BSTATE程序分析與重構(gòu)進展報告

3.1 編譯系統(tǒng)

重構(gòu)前的BSTATE僅能采用GNU make進行程序構(gòu)建和管理.Makefile文件描述了整個工程的編譯、連接等規(guī)則.使用 Makefile 時, BSTATE的編譯步驟是:

1)通過預(yù)編譯指令(GNU的cpp或者intel的fpp), 生成編譯時使用的Fortran代碼.值得注意的是, BSTATE 的 mlwf(Maximally Localised Wannier Functions)部分需要使用自由格式的預(yù)編譯; 其余需要使用固定格式的預(yù)編譯.此外, 預(yù)編譯后實際生成的代碼會受條件編譯的控制.

2)預(yù)編譯后的源代碼 (.f90, .for)被統(tǒng)一為.f后綴的Fortran代碼.BSTATE編譯的過程中首先會編譯整個程序所需要的模塊函數(shù)(module),之后繼續(xù)編譯和連接所有的程序文件.

針對不同的計算機構(gòu)架和編譯環(huán)境, BSTATE有著不同的 Makefile模板.并支持 Intel MKL 和FFTW2等一系列的函數(shù)庫.針對不同的編譯環(huán)境,用戶需要重新修改所有的編譯選項.這對于非計算機專業(yè)的用戶存在不低的使用門檻, 同時使得跨平臺和交叉編譯復(fù)雜化.

重構(gòu)后的BSTATE在保留舊版Makefile編譯流程的基礎(chǔ)上, 提供全新的CMake編譯系統(tǒng)支持[25].CMake 全稱是“cross platform make”, 是一個開源的跨平臺自動化構(gòu)建系統(tǒng).新版的BSTATE使用指定名為CMakeLists.txt的配置文件來控制軟件的預(yù)編譯、編譯、連接等流程.與此同時,CMake天生跨平臺的特征, 使得僅需要對于CMakeList.txt進行簡單配置, 原則上即可以生成不同目標平臺 (如 Linux, Windows, Mac)的程序文件.這可以大大簡化跨平臺和交叉編譯方面的工作.重構(gòu)后的BSTATE編譯系統(tǒng)的特點主要有:

1)保留原版Makefile的編譯流程.使用cmake或者ccmake時可以指定預(yù)編譯環(huán)境(cpp、fpp).預(yù)編譯時可以指定條件編譯的選項, 生成編譯BSTATE時使用的有效代碼.

2)內(nèi)置了不同優(yōu)化等級的預(yù)編譯和編譯方案(如 debug, release 等); 支持多線程編譯.

3)提供Libxc、FFTW2/3、CUFFT/CUFFTW等外接函數(shù)庫開關(guān)選項; 支持不同語言的混合編譯.

4)通過 cmake-curses-gui提供圖形交互式(GUI)的預(yù)編譯和編譯支持.

5)通過對“MATH_ROOT”或類似關(guān)鍵環(huán)境變量的查找, 提供對于常見數(shù)學(xué)函數(shù)庫的自動搜索和支持功能.該功能默認開啟, 并按照“MKL,ESSL, OPENBLAS, ATLAS, ACML, SYSTEM_NATIVE”的默認順序進行數(shù)學(xué)函數(shù)庫的自動配置.

通過以上的調(diào)整, 重構(gòu)后的BSTATE可以通過CMake實現(xiàn)自動和快速編譯.比如終端中鍵入FC = mpif90 ccmake [BSTATE_CMakeList.txt],即可開啟預(yù)編譯和編譯的GUI, 進行自動配置或按提示修改編譯選項.

在表1中我們也總結(jié)了BSTATE新老編譯系統(tǒng)的對比.可以看出重構(gòu)后的編譯系統(tǒng)相對于舊版, 無論是在代碼維護、擴展性、用戶使用等等方面均有長足的進步.

表1 重構(gòu)前后 BSTATE 編譯系統(tǒng)的對比Table 1.Comparison of BSTATE compilation system.

3.2 接口升級、函數(shù)庫整合

重構(gòu)前的BSTATE程序已經(jīng)可以調(diào)用多種函數(shù)庫來實現(xiàn)額外的功能或加速計算過程, 例如傅里葉變換FFTW函數(shù)庫, 但其僅支持老舊的FFTW2接口.因此, 首先在BSTATE里面增加FFTW3[26]的調(diào)用接口, 以便使用常用的SSE2和AVX指令集, 加速運算.同時, 也初步接入了基于 CUDA的傅里葉變換函數(shù)庫(CUFFTW和CUFFT), 已可取得同CPU版本一致的計算結(jié)果, 但是效率仍有很大的提升空間.

重構(gòu)前的 BSTATE已經(jīng)支持 LDA, GGA,LDA+U, GGA+U 等計算.但其僅通過固定的代碼支持 ldapw91和 ggapbe兩個泛函.因此將Libxc庫[27]接入BSTATE, 使其可以支持流行的各種泛函.Libxc是當前最為全面的交換關(guān)聯(lián)泛函庫 .該 泛 函 庫 由 ETSF(European Theoretical Spectroscopy Facility)維護, 并幾乎收錄了主流期刊公開報道的所有交換關(guān)聯(lián)泛函.

Libxc內(nèi)部函數(shù)相互依賴, 調(diào)用需要編譯為庫文件后才能在子程序中調(diào)用.因此重構(gòu)時主要通過預(yù)編譯的方式在多個子程序中配置, 然后再進行調(diào)用.在重構(gòu)的BSTATE程序中, LDA和GGA泛函通過新添加的extern_functionals.f90接口程序接入.對于LDA泛函, 勢能僅僅依賴于電子的密度 (即); 對于 GGA 泛函, 勢能會依賴于電子的密度和電子密度的梯度(即接口程序定位在BSTATE中的xc_gga.f90源程序中.除了電子密度、梯度等信息, k點信息、自旋等也要在接口中傳遞.

圖2 閃鋅礦結(jié)構(gòu)的 GaAs雙原子化合物Fig.2.GaAs compound with zinc blende structure.

作為測試, 使用砷化鎵體系(GaAs, 見圖2)和石墨烯片段作為測試的算例.測試的結(jié)果表明重構(gòu)后的BSTATE在DFT計算中性能可以與之前持平(表2和表3), 無論是單核、多核還是MPI并行計算, 性能均沒有明顯差距.此外, 受益于FFTW3對于SSE2和AVX系列指令集的支持,在Intel異構(gòu)眾核平臺(MIC)的計算結(jié)果會有一定的提升.

表2 Libxc 版本與原始版本性能比較Table 2.Benchmarks between BSTATEs with/without Libxc.

表3 FFTW3+Libxc 版本與原始 FFTW2 版本性能比較Table 3.Benchmarks between V.FFTW2 and V.FFTW3+ Libxc.

圖3 多種泛函計算的 GaAs 的態(tài)密度 (a) (b) (c) (d)Fig.3.Calculated density of states (DOS): (a); (b); (c); (d).

使用整合 Libxc和 FFTW3的 BSTATE版本, 我們在原有的lpw91(LDA)和ggapbe(GGA)泛函的基礎(chǔ)上使用相對論(lda_x_rel[28])和自適應(yīng)(lda_x_rae[29])的LDA泛函計算了GaAs的態(tài)密度.其計算結(jié)果列于圖3, 可以看出對于該體系不同的泛函會得到整體相似的態(tài)密度結(jié)果, 主要區(qū)別在–6 eV、0—4 eV 能量區(qū)間.

3.3 并行分析和測試

提高第一性原理計算模擬軟件的計算效率一直是計算科學(xué)領(lǐng)域的一大挑戰(zhàn).其主要原因是電子的強相互作用導(dǎo)致能量和波函數(shù)只能在哈密頓構(gòu)造完之后才可以求解; 而且構(gòu)造和求解的過程中也涉及各種數(shù)據(jù)結(jié)構(gòu)或者數(shù)值代數(shù)方面的問題.BSTATE軟件在設(shè)計之初就已經(jīng)很好地考慮這些方面的問題, 并已經(jīng)可以通過MPI模塊按照k點進行分布式求解.所謂 k點并行化, 即將k點均勻分配給 (并行環(huán)境中的)多個 CPU, CPU 之間的數(shù)據(jù)交換較少, 計算效率高.

為了詳細分析BSTATE的并行性能, 我們選用GaAs晶體和石墨烯片段分別進行小規(guī)模和大規(guī)模的測試計算.對于較小體系如GaAs晶體, 使用個人PC在多核(4線程)計算, 分析過程中的各個大類模塊所占時間比例統(tǒng)計于圖4中(編譯采用-O2優(yōu)化).可以看出, 傅里葉變換時間占比接近58%, 緊隨其后為特征值求解, 時間占比約為12%左右, 態(tài)密度、球貝塞爾函數(shù)計算部分合計占比均為15%左右.對于較大體系如石墨烯片段的計算(圖4), 我們使用56核MPI并行的集群進行了測試, 測試結(jié)果表明特征值求解部分占用最多的時間比例, 約為65%, 同時傅里葉變換時間占比下降到18%左右.

圖4 個 人 PC (左 圖)及 集 群 MPI (右 圖)測 試 中BSTATE各模塊的時間占比統(tǒng)計Fig.4.Time percentages of subroutines in BSTATE using PC (left) and Cluster (right).

圖5 BSTATE 的 MPI多 線 程 并 行 分 析 ; 其 中 藍 色 、 紅 色 、 黃 色 、 黑 色 分 別 代 表 application(BSTATE), MPI_barrier,MPI_reduce以及MPI_allreduce.Fig.5.Analysis via TAU performance system framework for BSTATE.

在圖4中集群MPI的并行模式下面, 我們也通過使用 TAU (Tuning and Analysis Utilities;重構(gòu)的BSTATE通過CMake可以整合TAU模塊) 分析了BSTATE并行時各個線程的執(zhí)行情況,并示于圖5中.可以看出在按照現(xiàn)有的基于k點的并行策略, BSTATE各個線程有著較好的負載均衡, 各個計算核心的計算能力都得到很好的發(fā)揮.

最后, 我們使用重構(gòu)的BSTATE版本(即包含 CMake+TAU+FFTW3+Libxc 的版本), 在中國科學(xué)院的“元”集群上對石墨烯片段體系進行了更大規(guī)模的測試, 以驗證新的版本可以穩(wěn)定的執(zhí)行計算任務(wù).該測試平臺基于曙光TC4600 E刀片平臺, 采用 Intel Xeon E5 2680 V3 處理器 (2.5 GHz、12 核).測試的結(jié)果示于圖6 中.可以看出, 重構(gòu)的BSTATE相對于重構(gòu)前基于FFTW2的版本在跨節(jié)點并行方面有著一定的性能提升(0—17%左右的加速).這應(yīng)該得益于新的FFTW3調(diào)用接口以及基于Libxc的泛函求解方案.同時我們也注意到當前基于k點的并行求解方案在核心過多的情況下效率下降很快, 這也是實際計算中特別要注意的方面.

圖6 重構(gòu)前后的 BSTATE 跨節(jié)點并行測試Fig.6.Benchmark of MPI parallel for BSTATE with and without refactoring.

4 結(jié) 論

本文匯報了我們在重構(gòu)第一性計算模擬軟件BSTATE的一些最新進展.根據(jù)重構(gòu)的核心思想(降低用戶使用門檻、擴展軟件適用范圍、提高軟件計算效率), 我們的工作主要集中在編譯、配置部分的自動化, 增加泛函類別, 以及優(yōu)化程序接口3個方面.重構(gòu)后的BSTATE通過CMake的引入可以實現(xiàn)跨平臺、簡單圖形化的編譯; 通過Libxc和FFTW3的引入和更新實現(xiàn)了程序功能的擴展和對更多計算平臺的支持; 最后的并行分析和穩(wěn)定性測試, 驗證了重構(gòu)后的BSTATE可以更好地勝任第一性原理方面的計算任務(wù), 并在跨節(jié)點并行方面有著一定的性能提升.

在重構(gòu)的進程中, 也注意到BSTATE仍有較大的改進空間.比如在原理層面上的哈密頓的構(gòu)造、子空間迭代等自洽迭代改進方案, 以及技術(shù)層面上的CPU-GPU協(xié)同計算方面.這些需要更多的第一性原理算法上的積累和深入, 也是我們后面需要進一步考慮的方面.

感謝徐順博士關(guān)于程序優(yōu)化方面的指導(dǎo)和討論.