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

(武漢工程大學(xué), 光學(xué)信息與模式識(shí)別湖北省重點(diǎn)實(shí)驗(yàn)室, 武漢 430205)

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

1 引 言

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

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

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

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

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

2 理論方法

2.1 三維含時(shí)薛定諤方程求解方案

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

2.2 并行計(jì)算方案

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

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

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

3 結(jié)果與討論

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

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

3.2 實(shí)驗(yàn)結(jié)果

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

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

表2 不同角量子數(shù)下CPU 與GPU 的計(jì)算時(shí)間比較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)格點(diǎn)R的大小, 然后將CPU 與GPU 的計(jì)算時(shí)間對(duì)比,如表3 和圖3 所示.

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

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

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

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

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

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

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

動(dòng)能譜的誤差計(jì)算公式為

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

表5 不同矩陣大小下CPU 與GPU 的計(jì)算時(shí)間比較 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)動(dòng)量分布(a)CPU計(jì)算結(jié)果;(b)GPU計(jì)算結(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 實(shí)驗(yàn)結(jié)果討論

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

4 結(jié) 論

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