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

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

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

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

1 引 言

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

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

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

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

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

2 理論方法

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

上述對應算法步驟如表1 所列.

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

連續態波函數在球諧函數表象下可以寫成

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

電子與電場相互作用后的末態動量分布是

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

2.2 并行計算方案

本文在Matlab 環境下調用設備端(GPU)來實現并行計算, 離散化的初始波函數以及傳播算符均保存在主機端(CPU), 如圖1 所示. GPU 上數組的創建和傳輸通過Matlab 并行計算工具箱的相關函數完成, 使用gpuArray()函數從主機端向設備端發送Matlab 數組, 即將Matlab 工作區內的數組傳輸到設備端內存. 在GPU 中進行數值計算完畢以后, 再通過gather()函數從設備端向主機端發送Matlab 數組, 即將設備端內的Matlab 數組傳輸到主機端內存.

圖1 數據傳輸流程圖Fig. 1. The flowchart of data transmission.

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

3 結果與討論

3.1 測試環境與測試算例

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

3.2 實驗結果

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

首先, 設置徑向網格點R= 1024 × 32 不變,選取不同的角量子數L, 計算結果如表2 和圖2 所示.

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

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

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

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

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

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

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

圖3 加速比隨著徑向網格點的變化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.

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

動能譜的誤差計算公式為

圖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 氫原子的光電子末態動量分布(a)CPU計算結果;(b)GPU計算結果Fig.6.Photoelectron final-state mom entum distributions of hyd rogen atom:(a)Calculation results of CPU;(b)calculation results of GPU.

3.3 實驗結果討論

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

4 結 論

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