陳富州 程晨 2) 羅洪剛 2)?

1)(蘭州大學物理科學與技術學院,蘭州 730000)

2)(北京計算科學研究中心,北京 100084)

1 引 言

密度矩陣重正化群方法(DMRG)[1,2]是研究相互作用量子系統(tǒng)最重要的多體數值方法之一.眾所周知,多體量子系統(tǒng)的復雜度隨著系統(tǒng)尺寸指數增長,給數值計算帶來了極大的困難.而DMRG給出了一種非常有效的截斷希爾伯特空間的方法,保留有限個狀態(tài)并通過掃描即可得到變分收斂的基態(tài)或低激發(fā)性質.應用于自旋1的海森伯鏈時,該方法僅保留數百狀態(tài)并通過數次掃描即可得到相對誤差 1 0-9左右的基態(tài)能量,而所需的計算量僅隨格點尺寸線性增長[3,4].作為求解一維格點模型基態(tài)的成熟方法,DMRG也不斷被應用于其他各類問題并取得了一定的成功,如動量空間中的哈密頓量[5]與量子化學問題[6-8]、量子系統(tǒng)的時間演化問題[9-11]、準二維以及二維量子格點模型[12-14]等.這些嘗試極大擴展了DMRG的應用,但同時也對該數值算法的優(yōu)化提出了更高的要求.以二維相互作用電子系統(tǒng)為例,該系統(tǒng)包含非常豐富的物理,例如高溫超導條紋相[15-17]、量子自旋液體[18,19]等.然而,在面對二維或者準二維格點模型時,DMRG所需的保留狀態(tài)數大致隨格點寬度指數增長,得到收斂結果所需的掃描次數也大大增加.此時,DMRG所需的計算量和存儲量較大,使用各種方法優(yōu)化算法顯得十分必要.

在實際應用中,對DMRG算法的優(yōu)化與加速主要體現在兩個方面.一方面,人們不斷改進DMRG算法本身以減少計算量,包括使用各種好量子數對角化小的希爾伯特子空間[20,21],使用動態(tài)保留狀態(tài)數節(jié)約計算資源[22,23],使用好的初始預測波函數減少對角化方法迭代次數[24],使用單格點DMRG方法減少計算量與內存[25,26]等.另一方面,人們結合DMRG算法的特性,發(fā)揮高性能計算機的并行計算能力進一步縮短計算時間,包括實空間并行[27]、共享存儲的多核行[28]、分布式存儲的多節(jié)點行[29]以及CPU-GPU異構并行[30]等.

相比CPU,GPU具有較高的浮點計算能力以及較大的存儲帶寬,因此在通用計算中得到了大量的應用.當前很多求解強關聯格點模型的數值方法(比如：精確對角化方法[31]、量子蒙特卡羅方法[32]以及張量乘積態(tài)方法[33]等)都實現了基于GPU的并行優(yōu)化.在CPU-GPU異構并行環(huán)境中,異構并行算法可以同時發(fā)揮CPU和GPU的計算性能,并且可以實現內存和GPU顯存的分布式存儲,在一定程度上減小了GPU顯存容量對計算規(guī)模的限制.最近,Nemes 等[30]提出了DMRG方法的一種CPU-GPU異構并行實現,并將其應用到一維格點模型的基態(tài)能量計算.在保留狀態(tài)較多時,最耗時的哈密頓量對角化部分在GPU中獲得了接近峰值的運算性能.在對角化過程中,對占用存儲空間較多的Davidson方法實現了內存和GPU顯存的分布式存儲,減小了GPU顯存的容量限制.然而與此同時,Davidson方法在計算中內存和GPU顯存之間需要通信向量數據,其性能會受到通信帶寬的限制.另一方面,該方案中哈密頓量和波函數基于兩子塊表示,其子塊算符需要的存儲量大致相當于目前流行的四子塊表示的d2倍(d為單格點希爾伯特空間維數),在實際應用中存在很大的局限性.

在前人工作的基礎上,本文提出了一種新的CPU-GPU異構并行優(yōu)化算法,并使用四子塊表示的DMRG超塊與波函數.為了說明該方法的有效性,將其應用到準二維模型的求解,獲得了不同保留狀態(tài)數下的基態(tài)能量及其性能基準.在Hubbard梯子模型的例子中,得到了非均勻的電荷密度分布,與高溫超導銅氧面中觀測到的條紋相定性一致.我們希望CPU-GPU異構并行算法能進一步推進DMRG在求解二維與準二維模型、時間演化、量子化學等問題中的應用,同時引起強關聯領域對GPU算法的關注和重視.

本文的內容包括以下幾部分：第2節(jié)回顧了有限DMRG算法,基于四個子塊表示的DMRG實現,介紹了該工作采用的基準模型,并得到了在CPU執(zhí)行時各個部分的計算時間占總時間的比例;第3節(jié)針對DMRG中最耗時的哈密頓量對角化部分給出了異構并行優(yōu)化方法;第4節(jié)以計算4腿Hubbard模型基態(tài)為例,對比了異構并行優(yōu)化方法與單個CPU中MKL并行的性能;最后給出本文的總結.

2 有限DMRG方法和基準模型

有限尺寸格點模型的DMRG計算分為兩部分：首先執(zhí)行無限DMRG方法,這可為后面的計算提供較好的初始波函數;然后執(zhí)行有限DMRG掃描,其中每一步有限DMRG掃描會優(yōu)化系統(tǒng)塊的狀態(tài),直到收斂至基態(tài).在準二維梯子模型的計算中,有限DMRG部分需要多次掃描,且保留很多的狀態(tài)才能收斂,幾乎占用整個DMRG計算的全部時間,因此本文主要考慮該部分的并行優(yōu)化.有限DMRG計算中向左掃描和向右掃描非常相似,本文以向右掃描為例介紹有限DMRG方法.每一步優(yōu)化的過程中整個格點系統(tǒng)由四個部分構成 (圖 1)：S,s1,e1 和E,其中S和s1 構成系統(tǒng)塊,E和e1 構成環(huán)境塊,系統(tǒng)塊和環(huán)境塊構成超塊.基于四個子塊,波函數有如下形式：

其中|s〉,|e〉,|σs1〉和|σe1〉分 別為子塊S,E,s1 和e1上的狀態(tài).同樣,基于四個子塊,一般形式的哈密頓量H可以表示為

其中OS,OE,Os1和Oe1分別為子塊S,E,s1 和e1上的算符.以求解系統(tǒng)的基態(tài)為例,給出基于四子塊表示的DMRG一步優(yōu)化過程,如算法1所示.重復執(zhí)行DMRG的一步優(yōu)化過程,進行多次實空間掃描,即可得到收斂的結果.在算法1中,對角化超塊哈密頓量最為耗時,通常人們采用稀疏矩陣對角化方法(Lanczos方法、Davidson方法等).此時,不需要超塊哈密頓量的矩陣表示,僅需要算符在四個子塊中的矩陣表示,并迭代執(zhí)行|φ〉=H|ψ〉,文獻[4]中給出了該操作的高效算法,具體過程如算法2.在后面的描述中,我們稱算法2中第2,5和8行對應的循環(huán)分別為step1,step2和step3.可以看到該算法(st)ep1和step3中為矩陣乘法,對應的計算量為OD3,其中D為DMRG保留狀態(tài)(數.)step2中僅包含向量操作,對應的計算量為OD2.在保留狀態(tài)較多時,哈密頓量作用在波函數的計算量主要由step1和step3決定.而對角化哈密頓量的總時間線性依賴于對角化方法的迭代次數,本文實現的有限DMRG中通過上一步優(yōu)化波函數給出一個較好的初始迭代波函數,可以有效地加快對角化方法的收斂[24].

圖1 超塊中的四個子塊Fig.1.4 Sub-blocks of super-block.

?

?

為了獲得異構并行優(yōu)化方法的性能基準,選取4腿梯子上的Hubbard模型為例子,并利用DMRG求解其基態(tài).具體地,系統(tǒng)哈密頓量為

這里t為兩個格點之間的電子躍遷能,U為單個格點上的電子在位庫侖排斥能,為格點i上自旋σ的費米子產生(湮滅)算符,為自旋σ的粒子數算符,其中σ∈{↑,↓}.在梯子模型中,格點i包含x,y兩個方向的分量,以x標記梯子長邊坐標,y標記梯子短邊坐標.在后面的測試計算中,選取t=1 為能量單位,沿著長邊和短邊方向分別使用開邊界和周期邊界條件,梯子長度為16.對于其他參數,取相互作用U=8,總電荷密度0.875,即 1 /8 空穴摻雜,這是在其他數值工作中觀測到條紋相的典型參數[13,34].另外本文的實現利用了該模型的總粒子數和總自旋在z方向投影兩個好量子數,每個子塊的希爾伯特空間可以被劃分為多個子空間,每個子空間中的狀態(tài)對應于相同的量子數.此時模型(1)中算符的矩陣表示為分塊矩陣,相應地,超塊的希爾伯特空間可以用其四個子塊的好量子數劃分為多個子空間.

模型(1)哈密頓量的矩陣表示為實對稱矩陣,本文使用Davidson方法[35,36]對角化該哈密頓量.Davidson方法是一種使用預條件技術的子空間迭代方法,算法3給出了該方法每一步迭代的具體操作.可以看到,每一步迭代都需要作用哈密頓量到波函數,并且包含多個向量操作.其中向量操作的計算量和存儲量均線性依賴于Davidson方法子空間中向量的個數和向量的維數.在性能測試中,使用7次有限DMRG掃描收斂到基態(tài),Davidson方法子空間中向量個數最大為11.

?

本文所有計算使用的異構并行環(huán)境包含1個CPU (Intel Xeon E5-2650 v3,10核,主頻2.3 GHz)和Nvidia Tesla K80卡中的1個GPU (含12 GB顯存).程序的實現基于CUDA環(huán)境,CPU和GPU中的矩陣向量操作分別調用了Intel MKL和CUBLAS中的子程序.為了描述矩陣操作的性能,估計了單位時間內處理器執(zhí)行的浮點操作的次數,其中總的浮點操作次數在保留狀態(tài)數較大時可以由所有矩陣乘法中浮點操作次數之和近似.為獲取異構并行的優(yōu)化表現,需要給出CPU上的并行性能作為基準.哈密頓量對角化中包含大量矩陣乘法,這里首先測試了隨機方陣乘法的性能,結果見圖2(a).可以看出,當矩陣尺寸較大(大于400)時,矩陣乘法的性能較高,并隨著矩陣尺寸增大而增大.對于DMRG算法,圖2(b)中的結果表明,保留狀態(tài)數越大并行計算性能越高,并逐漸接近峰值性能.另外也測試了對角化哈密頓量以及哈密頓量作用于波函數部分占總計算時間的比例,如圖3 所示.在我們所關心的保留狀態(tài)數范圍內,對角化哈密頓量的耗時比例超過總計算時間的90%,其中作用哈密頓量到波函數占總時間比例超過80%.因此,我們的工作主要針對該部分進行異構并行優(yōu)化.

圖2 CPU中作用哈密頓量在波函數上的性能 (a)矩陣乘法的浮點性能;(b)作用哈密頓量于波函數的性能,及矩陣乘法中的最大矩陣尺寸Fig.2.Performance of acting the Hamiltonian on the wave function in CPU：(a)The matrix multiplication performance;(b)the performance of acting the Hamiltonian on the wave function,and the maximum matrix size of the matrix multiplications.

圖3 對角化哈密頓量和作用哈密頓量到波函數操作占總計算時間的比例Fig.3.Time ratio of diagonalization of the Hamiltonian and acting the Hamiltonian on the wave function to the total time cost.

3 DMRG的異構并行實現

主要考慮DMRG方法在準二維模型中的應用,并且針對有限DMRG中計算量較大的哈密頓量對角化部分進行并行優(yōu)化.對于準二維問題,DMRG方法達到較高精度通常需要保留較多的狀態(tài),相應地,對角化哈密頓量時會執(zhí)行一些大尺寸矩陣操作.類似于CPU,矩陣乘法的性能在GPU中隨著矩陣尺寸增大而增大;同時,GPU的浮點運算能力一般遠大于單個CPU.因此,在異構并行優(yōu)化中,我們傾向于盡可能在GPU中執(zhí)行大尺寸的矩陣操作.從存儲方面考慮,為了避免GPU顯存和內存之間頻繁的數據通信,多次參與GPU計算的數據需要存儲在GPU顯存中,主要包括Davidson方法中的向量、算符數據和臨時數據(算法2中但同時,保留較多的狀態(tài)也導致計算中需要的存儲容量較大;考慮到當前GPU顯存容量較小,在異構并行方法中需要合理分配各個部分的存儲利用.

首先考慮Davidson方法的異構并行實現,如算法3所示,Davidson方法中主要執(zhí)行各種向量操作,相對而言運算量較小,因此應盡可能充分發(fā)揮內存和GPU顯存的存儲帶寬.具體地,我們將所有向量以相同的方式按行劃分為兩部分,使得一部分波函數基矢對應的分量存儲在GPU顯存中,另一部分存儲在內存中.這樣任一向量的操作將由CPU和GPU共同完成,并且不需要向量數據的通信.當GPU顯存足夠時(通常內存容量遠大于GPU顯存),內存和GPU顯存中向量行數的比值為兩者存儲帶寬的比值,理論上此時可以獲得最高的性能.由于Davidson方法的存儲量線性依賴于子空間中向量的個數和超塊希爾伯特空間的維數,當向量個數或者DMRG保留狀態(tài)較大時,GPU顯存中的向量數據會受到GPU顯存容量的限制.這會導致較多向量操作在CPU中執(zhí)行,此時Davidson方法獲得的性能較低.

接下來給出哈密頓量作用于波函數(即算法2)部分的異構并行實現,首先將其中的操作合理劃分到CPU和GPU中.將算法2中所有操作按照子塊S中的好量子數分為多個組,此時各個組中的運算可以獨立進行,僅在求和計算(即算法2中step3)時需要通信.此時波函數〉也 按S中好量子數劃分為兩個部分,其中一部分僅在GPU中計算(記為另一部分僅在CPU中計算(記為

算法2中耗時較多的運算為一系列矩陣乘法,且S中每個好量子數對應分組的計算量在執(zhí)行運算前可以比較準確地估計.通常,GPU的浮點運算能力強于CPU,適合處理大尺寸矩陣乘法運算,因此在這一步盡量將大矩陣運算分配至GPU,將相對較小的矩陣運算分配至CPU執(zhí)行.為了實現這一目標,在具體操作中,我們將矩陣乘法平均運算量較大的組分配到GPU中執(zhí)行,這里平均運算量為組內矩陣乘法總計算量與矩陣乘法個數的比值.進一步,可根據GPU中計算量占的比例PGPU將相互獨立的分組計算分配給GPU,剩余組的計算分配給CPU執(zhí)行.

為了獲得較高的異構并行效率,在執(zhí)行Davidson對角化方法時動態(tài)調整PGPU以盡可能實現負載均衡.在具體操(作中,將)每一步迭代優(yōu)化比例所處的區(qū)間記為P0,P1.設定初始區(qū)間P0=0,P1=1,然后進入Davidson(方法迭代)過程.令GPU中計算量比例PGPU=P0+P1/2,執(zhí)行一步Davidison迭代,可以獲得此時作用哈密頓量到波函數的CPU和GPU計算時間,分別記為TCPU和TGPU.以此為依據更新下一步迭代區(qū)間,使得

并開始下一步Davidson迭代.如此按照類似二分法的收斂思路,通過少數幾步迭代就可以收斂到優(yōu)化比例,之后的大部分迭代運算都趨于負載均衡.

本文主要針對保留狀態(tài)數較大的問題進行優(yōu)化,此時涉及到的矩陣較大,GPU中矩陣運算的并行效率較高.同時,考慮到GPU顯存的限制,為了減小臨時數據的存儲量,根據子塊S分組后各個組的操作依次執(zhí)行.這種情況下,僅需要分配一段存儲空間,使其同時滿足任意一個組中的所有操作即可.在圖4中,給出了對角化哈密頓量部分運算的存儲需求,并給出了與兩子塊表示所需存儲的對比.可以明顯看出其總體的顯存需求遠遠小于兩子塊表示.因此,相比于參考文獻[30]中的異構并行,本文方案可以處理需要更大DMRG保留狀態(tài)數的問題.

圖4 存儲臨時數據,子塊算符需要的GPU顯存Fig.4.The GPU memory cost of temporary data and subblock operators.

4 數值結果

為了說明本文優(yōu)化方法的有效性,我們分別保留4096,6144,8192,10240個狀態(tài)計算4腿Hubbard梯子的基態(tài)能量,并得到了相應的性能基準.圖5(a)給出了DMRG優(yōu)化中各個部分相對于單個CPU的加速比(單個CPU計算時間與單個GPU計算時間的比值),可以看到作用哈密頓量在波函數部分獲得加速比最大,其加速比在保留狀態(tài)數大于4096后較為接近(不低于3.8).當保留狀態(tài)數較大時,由于GPU顯存總量的限制(如圖5(b)所示),大部分向量操作由CPU完成,這導致Davidson方法的加速比在保留狀態(tài)數較大時有所下降.然而,Davidson方法中向量操作占哈密頓量對角化的時間比例較少,因此對哈密頓量對角化部分加速比影響較小(不低于3.5).本文哈密頓量對角化之外的其他操作計算時間占總時間比例約10% (如圖3),該部分的GPU并行優(yōu)化還沒有被考慮,因此總并行加速比低于哈密頓量對角化部分的加速比,這里保留最大10240個狀態(tài)獲得的加速比為2.85.圖5(c)中給出了異構并行實現中CPU和GPU分別貢獻的浮點性能,可以看到隨著保留狀態(tài)數的增大,CPU和GPU中執(zhí)行的大尺寸矩陣增多,兩者貢獻的浮點性能隨之增大.雖然本文數值計算保留狀態(tài)數較大,但是其中矩陣尺寸為兩子塊表示時的 1 /d(Hubbard模型中d=4),目前獲得的性能仍明顯小于GPU可以達到的峰值性能(1200 GFLOPS).對 1 6×4 的梯子,根據不同保留狀態(tài)數(4096,6144,8192,10240)得到的基態(tài)能量外推得到模型(1)在U=8.0 時的格點平均能量Eg=-0.75114(2),與文獻[34]結果一致,見圖6.進一步給出基態(tài)的電荷密度分布,如圖7所示,可以觀察到明顯的電荷密度條紋,這是銅氧化物高溫超導體中經常被觀測到的現象之一[13,14,37,38].

圖5 異構并行的性能 (a)加速比;(b)Davidson方法中的向量占用GPU顯存;(c)作用哈密頓量到波函數部分的性能Fig.5.Performance of hybrid parallel strategy：(a)The speedup;(b)the GPU memory cost of vectors in Davidson;(c)the performance ofH|ψ〉.

圖6 基態(tài)能量關于截斷誤差的函數(直線表示對基態(tài)能量的線性外推,直至截斷誤差為0)Fig.6.Groundstate energy as a function of truncation error.The straight line gives a linear extrapolation of the ground energy until 0 truncation-error.

圖7 對于16×4 Hubbard模型,U=8.0時的基態(tài)電荷密度分布(可以觀察到明顯的電荷密度條紋)Fig.7.Ground state density profile for the 16×4 Hubbard ladder withU=8.0.Charge density stripes can be clearly observed.

5 結 論

本文主要考慮DMRG方法在準二維格點模型中的應用,針對其中最耗時的哈密頓量對角化部分實現了CPU-GPU異構并行優(yōu)化,并且給出了負載平衡方法.為了減小準二維格點模型計算中GPU顯存的限制,本文的實現中哈密頓量與波函數基于四子塊表示,其對角化時需要的GPU顯存占用遠小于兩子塊表示,使得本文的異構并行方法可以應用于更多模型、更多問題的研究.將該方法應用到4腿Hubbard梯子模型的求解中,得到了不同保留狀態(tài)數時DMRG中各個部分的加速比.數值結果表明,本文的異構并行方法適用于保留狀態(tài)數較大的準二維模型計算,并且總性能隨著保留狀態(tài)數增大而增大.目前,強關聯物理問題很大程度上依賴于多體數值計算,一些復雜問題通常進一步受制于計算方法的計算量與計算時間.在多體算法本身出現革命性發(fā)展之前,合理利用計算機技術的發(fā)展提升算法的效率能為研究強關聯系統(tǒng)提供很大的幫助.我們希望該并行方法可以在更多的復雜格點模型、更多問題中得到應用,并能夠進一步引起強關聯領域對于以GPU為代表的新技術的關注和重視.