基于圖形處理器加速數(shù)值求解三維含時薛定諤方程*

唐富明 劉凱 楊溢 屠倩 王鳳 王哲 廖青

(武漢工程大學, 光學信息與模式識別湖北省重點實驗室, 武漢 430205)

量子力學領(lǐng)域中對強激光場與原子分子相互作用的理論研究非常依賴于數(shù)值求解含時薛定諤方程. 本文在強場電離的背景下并行求解氫原子的三維含時薛定諤方程. 基于球極坐標系, 采用分裂算符-傅里葉變換方法將含時薛定諤方程進行了離散化. 由此可得到長度規(guī)范下的光電子連續(xù)態(tài)波函數(shù). 圖形處理器(GPU)可以依托多線程結(jié)構(gòu)充分發(fā)揮細粒度并行的優(yōu)勢, 實現(xiàn)整體算法的并行加速. 計算表明, 相對于中央處理器(CPU), GPU 并行計算有著最高約60 倍的加速比. 由此可見, 基于GPU 加速數(shù)值求解三維含時薛定諤方程能夠顯著縮短計算耗費的時間. 這一工作對利用GPU 快速求解三維含時薛定諤方程有著重要的指導(dǎo)意義.

1 引 言

強場激光物理是過去二十多年隨著超短強激光技術(shù)的發(fā)展而快速發(fā)展起來的前沿學科, 主要研究超強超短激光與電子的相互作用. 早在1917 年,愛因斯坦的輻射理論就提出了受激輻射的基本概念, 預(yù)測到光可以產(chǎn)生受激輻射. 直至1960 年, 世界上第一臺激光器誕生, 此后的幾十年間激光技術(shù)飛速發(fā)展, 從自由輸出到調(diào)Q技術(shù)(Q-switch)、鎖模技術(shù)(mode locking), 再到啁啾脈沖放大技術(shù),激光的脈沖寬度越來越短, 功率越來越大. 利用先進激光技術(shù)獲得的超快強激光脈沖與物質(zhì)相互作用, 成為了研究物質(zhì)基本性質(zhì)的一種重要手段. 其中, 強場中一些非線性現(xiàn)象, 如高次諧波的產(chǎn)生(HHG)[1?5]、次序和非次序雙電離[6?8]等, 受到了廣泛的關(guān)注.

強場激光物理的主要理論方法是數(shù)值求解含時薛定諤方程(TDSE)[9?11]、強場近似(SFA)[12?14]和半經(jīng)典模型[15?17]. 求解三維含時薛定諤方程得到的結(jié)果可以認為相當于發(fā)生在數(shù)值上的實驗. 但是求解三維含時薛定方程并非一項簡單的任務(wù), 無法獲得解析解, 只能借助計算機數(shù)值求解. 在直角坐標系下哈密頓算符的表示很直觀, 但由于正交網(wǎng)格在三個維度上均較為致密, 因此對存儲空間和計算量的需求十分巨大. 而在球極坐標系中只有徑向網(wǎng)格較為致密, 另外兩個維度相對稀疏. 盡管如此,求解三維含時薛定諤方程的計算量也是十分巨大.為了縮短計算的時間, 更快地得到計算結(jié)果, 就有必要使用并行方法去加速計算.

在過去的一段時間里, 對基于圖形處理器(GPU)并行求解三維含時薛定諤方程的優(yōu)化與加速主要體現(xiàn)在兩個方面. 一方面, 不斷改進TDSE算法或者引入其他方法以減小計算量, 如在柱坐標系下采用的混合數(shù)值迭代格式[18]. 另一方面, 采用不同的并行平臺, 比如: 使用CUDA (compute unified device architecture)平臺編寫TDSE 程序[19], 使用OpenCL(open computing language)語言編寫TDSE 程序[20?22]. 但在CUDA 架構(gòu)或OpenCL 架構(gòu)下對算法并行計算進行加速, 使用起來難度較高.

基于硬件的不同, 并行加速有中央處理器(CPU)并行和GPU 并行兩種方法. CPU 中多條指令構(gòu)成指令流水線, 且每個線程都有獨立的硬件來操縱整個指令流. 采用復(fù)雜的分支預(yù)測技術(shù)來達到并行計算目的. GPU 是圖形處理器, 圖形運算的特點是大量同類型數(shù)據(jù)的密集運算—如圖形數(shù)據(jù)的矩陣運算, 正因為如此, GPU 的微架構(gòu)就是面向適合于矩陣類型的數(shù)值計算而設(shè)計的, 這類計算可以分成眾多獨立的數(shù)值計算—大量數(shù)值運算的線程, 而且數(shù)據(jù)之間沒有像程序執(zhí)行的那種邏輯關(guān)聯(lián)性. GPU 并行計算的崛起得益于大數(shù)據(jù)時代的到來, 而傳統(tǒng)的CPU 并行計算已經(jīng)遠遠不能滿足大數(shù)據(jù)的需求. CPU 的核處理器相對較少, 并行的效率較低. GPU 最大的特點是擁有超多計算核心, 往往成千上萬核. 而每個核心都可以模擬一個CPU 核心的計算功能. 因此GPU 的并行效率比CPU并行效率高.

本文基于GPU 加速數(shù)值求解三維含時薛定諤方程, 通過在GPU 上并行加速求解三維含時薛定諤方程, 并且與CPU 并行加速做對比, 得到GPU相對于CPU 的加速比. 在GPU 加速的應(yīng)用程序中, 工作負載的順序部分運行在CPU 上(這是為單線程性能優(yōu)化的), 而應(yīng)用程序的計算密集型部分并行運行在數(shù)千個GPU 內(nèi)核上. 正是如此, GPU并行計算可以得到很好的加速效果.

2 理論方法

2.1 三維含時薛定諤方程求解方案

物理系統(tǒng)的時間演化由波函數(shù)Ψ(t) 來描述, 滿足含時薛定諤方程

本文以線偏振激光電場作用下的氫原子為例, 在球極坐標中系統(tǒng)的哈密頓量為(如無特殊說明, 下文一律使用原子單位):

這里H(t) 指作用在約化波函數(shù)Φ(t)=rΨ(t) 上的算符, 其中L2為系統(tǒng)的總角動量算符,E(t) 為線偏振入射激光電場.

將波函數(shù)Ψ(r,t) 在球諧函數(shù) Ylm(θ,φ) 下展開,

由于m是一個好量子數(shù), 通過對方位角φ進行積分并把角動量算符數(shù)值求解.含時薛定諤方程的一般步驟是先給定一個初始波函數(shù)Φ(t=0) , 然后將時間演化算符作用在該初始波函數(shù)上, 反復(fù)迭代直至得到任意時刻的末態(tài)波函數(shù). 假設(shè)將波函數(shù)向前推進 Δt時間增量, 那么t時刻和t+Δt時刻的波函數(shù)關(guān)系為

利用分裂算符法[23]將(2)式代入(4)式, 可以得到

這里產(chǎn)生的變換誤差為 Δt的三階誤差項. 由于m是一個好量子數(shù), 通過對方位角φ進行積分并把角動量算符Lz替換為它的本征值m, 可以使問題大為簡化. 為了方便起見, 下文只討論m=0 的情況. 假設(shè)波函數(shù)可以在有限勒讓德多項式下展開:

Pl為球諧函數(shù) Yl0,fl(ri,t) 可以由高斯-勒讓德求積法數(shù)值計算得出,

這里xj為勒讓德多項式PL+1(xj) 的L+1 個零點,wj為與之對應(yīng)的求積權(quán)重. 對于l的不同取值, 一維徑向波函數(shù)fl(ri,t) 均定義在R點等距徑向網(wǎng)格上, 它的網(wǎng)格表示為

(4)式中時間演化算符的作用可以分為三步.

1) 徑向動能算符T=(-1/2)?2/?r2獨立作用在每個一維徑向波函數(shù)fl(ri,t) 上, 利用帶邊界條件的傅里葉譜方法[24], 得到

上述操作中, 由于對所有徑向波函數(shù)都進行了一維正弦變換, 應(yīng)用徑向動能算符的計算復(fù)雜度為O(Rlog2(R)L).

2) 等效勢能算符V分為離心勢能項L2/(2r2)和庫侖勢能項-1/r, 其中總角動量算符L2在球諧函數(shù)表象下是對角化的, 對角元為它的本征值l(l+1), 因此可以得到

3) 相互作用勢能算符W(r,θ,t) 在坐標表象{r,θ}下是對角化的, 利用 (5)式重構(gòu)波函數(shù)Φ(r,x,t) ,得到

上述對應(yīng)算法步驟如表1 所列.

通過(6)式再次展開成徑向波函數(shù)fl(ri,t) , 接下來只需要依次執(zhí)行步驟2)和步驟1), 就完成了波函數(shù)向前推進 Δt的一次迭代. 在含時演化中, 重復(fù)執(zhí)行上述步驟, 能夠得到任意時刻的末態(tài)波函數(shù). 步驟1)中的快速傅里葉變換操作以及(5)式和(6)式中的變換操作都可以通過向量化提高計算效率.

連續(xù)態(tài)波函數(shù)在球諧函數(shù)表象下可以寫成

其中是σl=argΓ(1+l-iZ/k) 是庫侖散射相移,δl是除去長程庫侖勢的短程勢產(chǎn)生的相移.

電子與電場相互作用后的末態(tài)動量分布是

表1 TDSE 算法步驟Table 1. TDSE algorithm steps.

2.2 并行計算方案

本文在Matlab 環(huán)境下調(diào)用設(shè)備端(GPU)來實現(xiàn)并行計算, 離散化的初始波函數(shù)以及傳播算符均保存在主機端(CPU), 如圖1 所示. GPU 上數(shù)組的創(chuàng)建和傳輸通過Matlab 并行計算工具箱的相關(guān)函數(shù)完成, 使用gpuArray()函數(shù)從主機端向設(shè)備端發(fā)送Matlab 數(shù)組, 即將Matlab 工作區(qū)內(nèi)的數(shù)組傳輸?shù)皆O(shè)備端內(nèi)存. 在GPU 中進行數(shù)值計算完畢以后, 再通過gather()函數(shù)從設(shè)備端向主機端發(fā)送Matlab 數(shù)組, 即將設(shè)備端內(nèi)的Matlab 數(shù)組傳輸?shù)街鳈C端內(nèi)存.

圖1 數(shù)據(jù)傳輸流程圖Fig. 1. The flowchart of data transmission.

在求解含時薛定諤方程時, 具體需要并行計算的部分是: 快速傅里葉變換、逆快速傅里葉變換、矩陣乘法和數(shù)組點乘. 將需要進行變換的數(shù)組從CPU 傳輸?shù)紾PU, 然后在GPU 上對數(shù)組按列分別進行快速傅里葉變換. 不同的列在進行快速傅里葉變換時是相互獨立的, 所以可以在不同的線程內(nèi)對不同的列同時進行快速傅里葉變換, 這樣就達到了并行的目的. 計算得到的數(shù)組存儲在GPU 內(nèi)存中, 將計算結(jié)果傳輸回CPU 后, 可以釋放在GPU中占用的內(nèi)存. 逆快速傅里葉變換與快速傅立葉變換類似. 矩陣乘法的思路是將左乘矩陣的任一行和右乘矩陣的任一列做點積, 得到目標矩陣的任一元素. 每個目標矩陣元素的計算都是相互獨立的, 因此可以同時計算多個目標矩陣元素, 以達到并行計算的目的. 數(shù)組點乘是Matlab 中的一種運算方式,直接對兩個數(shù)組中相同位置的元素做乘法, 即可得到目標數(shù)組的對應(yīng)元素. 因此, 數(shù)組點乘本質(zhì)上是一種標量運算, 非常便于利用多線程進行并行計算.

3 結(jié)果與討論

3.1 測試環(huán)境與測試算例

本文所有計算使用的并行環(huán)境包含1 個CPU(Intel Xeon E7-8880 v4, 22 核, 主頻2.2 GHz)和NVIDIA Tesla P100 卡 中 的1 個GPU (3584 核心, 主頻1.3 GHz). 程序的實現(xiàn)基于Matlab 環(huán)境,GPU 中的矩陣向量操作調(diào)用了重載的Matlab 函數(shù). 本次測試算例以線偏振紅外激光電場作用下氫原子的閾上電離為背景. 紅外激光電場強度I=9×1013W/cm2, 波長λ=800 nm , 具有正弦平方脈沖包絡(luò), 共8 個光學周期. 其中計算采用的空間范圍rmax=1000 a.u.,時間步長 Δt=0.0037 a.u..

3.2 實驗結(jié)果

為獲取GPU 并行的優(yōu)化表現(xiàn), 需要給出CPU上的并行性能作為基準. 通過改變角量子數(shù)L和徑向網(wǎng)格點R計算演化時間結(jié)束后的末態(tài)波函數(shù),比較不同參數(shù)下GPU 相對于CPU 的加速比(單個CPU 計算時間與單個GPU 計算時間的比值):

首先, 設(shè)置徑向網(wǎng)格點R= 1024 × 32 不變,選取不同的角量子數(shù)L, 計算結(jié)果如表2 和圖2 所示.

表2 不同角量子數(shù)下CPU 與GPU 的計算時間比較Table 2. Computation time of CPU and GPU under different angular quantum numbers.

圖2 加速比隨著角量子數(shù)的變化Fig. 2. Speedup ratio as a function of angular quantum number.

角量子數(shù)L= 19 保持不變, 改變徑向網(wǎng)格點R的大小, 然后將CPU 與GPU 的計算時間對比,如表3 和圖3 所示.

同時改變徑向網(wǎng)格點和角量子數(shù)的大小, 也就是改變徑向波函數(shù)集{fl(ri,t)}的矩陣大小, 比較CPU 和GPU 的計算時間, 結(jié)果如表4 和圖4 所示.

此外, 為了進一步測試GPU 的加速性能, 又進行了一組1600 nm 中紅外激光的實驗, 中紅外激光電場強度I=1×1013W/cm2,波長λ=1600 nm ,具有正弦平方脈沖包絡(luò), 共8 個光學周期. 其中計算采用的空間范圍rmax=1000 a.u., 時間步長Δt=0.0037 a.u.. 此時與第三組實驗相同, 同時改變徑向網(wǎng)格點和角量子數(shù)的大小, 也就是改變徑向波函數(shù)集{fl(ri,t)}的矩陣大小, 比較CPU 和GPU 的計算時間, 結(jié)果如表5 和圖5 所示.

選取一組800 nm 激光電場作用下氫原子的閾上電離, 由CPU 計算得到的末態(tài)波函數(shù)與GPU 的計算結(jié)果對比, 如圖6 所示, 電子末態(tài)動量分布保持一致.

表3 不同徑向網(wǎng)格點下CPU 與GPU 的計算時間比較Table 3. Computation time of CPU and GPU under different radial grid points.

圖3 加速比隨著徑向網(wǎng)格點的變化Fig. 3. Speedup ratio as a function of radial grid point.

表4 不同矩陣大小下CPU 與GPU 的計算時間 比較Table 4. Computation time of CPU and GPU under different matrix sizes.

同時計算兩種方法得到的約化波函數(shù)的誤差為 5×10?21,動量分布誤差為 2×10?18. 其中波函數(shù)誤差計算公式為

動能譜的誤差計算公式為

圖4 加速比隨著矩陣大小的變化Fig. 4. Speedup ratio as a function of the size of matrix.

表5 不同矩陣大小下CPU 與GPU 的計算時間比較 Table 5. Computation time of CPU and GPU under different matrix sizes.

圖5 加速比隨著矩陣大小的變化Fig. 5. Speedup ratio as a function of the size of matrix.

圖6 氫原子的光電子末態(tài)動量分布(a)CPU計算結(jié)果;(b)GPU計算結(jié)果Fig.6.Photoelectron final-state mom entum distributions of hyd rogen atom:(a)Calculation results of CPU;(b)calculation results of GPU.

3.3 實驗結(jié)果討論

通過上述的計算可以得出:1)相同計算量下CPU與GPU的計算時間有很大的差距,當計算量較小時加速比急劇升高,隨著計算量的增大加速比趨于一個穩(wěn)定值,最高達到了約60倍的加速提升,加速效果十分明顯;2)從不同的計算量對比也可以看出,當計算量越大時,加速的效果也就越好.并不是所有的計算都要選取GPU來計算,CPU將數(shù)據(jù)傳輸?shù)紾PU需要一定的時間,當數(shù)據(jù)比較大時,采用GPU來并行計算,能夠獲得更大的加速比;3)如果只是改變一個維度的大小,所得到的實際加速效果有時候并不理想,這和CPU以及GPU存儲數(shù)據(jù)、讀取數(shù)據(jù)的方式有關(guān).所以在計算的時候為了獲得最好的加速比,需要同時在至少兩個維度上調(diào)整矩陣的大小;4)無論采用CPU計算還是GPU計算,都能得到相同的計算結(jié)果,計算誤差也在可接受的范圍內(nèi),并且該結(jié)果也符合現(xiàn)有的閾上電離物理規(guī)律.

4 結(jié) 論

本文詳細分析了在強場電離的背景下數(shù)值求解氫原子三維含時薛定諤方程基于不同硬件的并行速度.借助于分裂算符-傅里葉變換方法,在球極坐標下得到了三維含時薛定諤方程的末態(tài)解.同時,依托于GPU的多線程結(jié)構(gòu),使得GPU發(fā)揮細粒度方面的并行優(yōu)勢,實現(xiàn)整體算法的并行加速.采用了CPU并行和GPU并行兩種加速計算模式,探討了兩者并行加速的性能.通過與現(xiàn)有的物理規(guī)律相對比,驗證了程序的正確性.計算結(jié)果表明,當計算量較小時GPU相對于CPU的加速效果不突出,隨著計算量的增大加速比迅速增加,然后趨于一個穩(wěn)定值.GPU并行對于數(shù)值求解三維含時薛定諤方程有著相對于CPU最高約60倍的加速.可見,計算量越大,采用GPU并行獲得的加速比越大.這一工作對利用GPU高效數(shù)值求解含時薛定諤方程有著重要的指導(dǎo)意義.