從格點量子色動力學應用看國產超算環(huán)境的基礎軟件

宮明，蔣翔宇，陳瑩，劉朝峰

1. 中國科學院高能物理研究所，北京 100049；2. 中國科學院大學，北京 100049

1 格點量子色動力學及其應用特征

人類所處的物質世界，從微觀粒子到整個宇宙，是按怎樣的規(guī)律運行的？答案的最新版本自20世紀中期以來逐漸成形，目前其被稱作“標準模型”。這個標準模型把萬物的相互作用歸結為4種：用廣義相對論描述的引力，其支撐了宇宙學模型；用量子規(guī)范場論描述的強相互作用、弱相互作用、電磁相互作用，其支撐了粒子物理標準模型。

為了找到更深刻的答案，對標準模型的檢驗工作也在持續(xù)進行中。目前絕大多數實驗結果與理論計算符合得相當好，但也有一些不確定性，這些不尋常的跡象有可能是新物理理論的濫觴之處。比如最近的一個熱點——關于繆子的反常磁矩的測量，實驗結果與理論計算似有分歧[1]。然而，微妙之處在于，其中的強相互作用貢獻的計算是非常困難的，需要采用超級計算機進行大規(guī)模數值模擬才有可能理清結果的系統(tǒng)誤差和統(tǒng)計誤差。這樣的計算采用的理論和方法就是格點量子色動力學（lattice quantum chromo dynamics，LQCD）。

一般說來，電磁相互作用和弱相互作用可以采用“微擾”的方法把計算逐級展開逼近，可以達到很高的精度。但這種方法對于強相互作用，尤其在較低的能量標度下，是徹底失效的。因此，基于計算機數值模擬的LQCD幾乎成為唯一有效的方法。在人類尋找超出標準模型的新物理理論的征途上，它的重要性不言而喻。這些年來，在國際頂級超級計算機上進行的自然科學計算中，LQCD往往是占用計算資源最多的應用。

LQCD把時空切分成四維的格子，把閔可夫斯基時空上的量子場論問題轉化為歐幾里得時空上的統(tǒng)計問題。在這里，時間和空間是等同對待的，而且都是分立且有限的，可以用外推的辦法把有限的計算推廣到無限的自然世界上。把代表夸克的（偽）費米子場定義在時空格點上，把代表膠子的規(guī)范場定義在相鄰格點之間的連線上。每個格點上的費米子場可以用12個復數表示，每條連線上的規(guī)范場可以用3×3的復矩陣表示，這樣的模型就可以在計算機上用內存對象來實現(xiàn)了。

LQCD的計算流程是：先產生規(guī)范場組態(tài)，再基于這些組態(tài)計算費米子的傳播子，然后用這些傳播子構造各種關聯(lián)函數，最后從這些關聯(lián)函數中擬合抽取目標物理量。這個過程涉及幾十種復雜的算法，計算規(guī)模和中間數據非常龐大，耗時常以年計，數據文件規(guī)模常以PB計。然而，這些不同算法的熱點函數頗為集中，數據雖大但非常規(guī)整，這為軟件設計和優(yōu)化提供了重要的有利條件。

LQCD的主要熱點函數可分為兩類：以Dslash函數為核心的大規(guī)模稀疏矩陣求解、以關聯(lián)函數縮并為代表的大規(guī)模張量計算。

稀疏矩陣求解的典型階數為108~109數量級，由于物理上的近鄰相互作用的特點，矩陣往往是稀疏的（也有不稀疏的情況，可以用稀疏矩陣的分式級數來逼近）。在其他研究領域的一般情況的稀疏矩陣往往是不規(guī)則的，需要在運行時動態(tài)尋址，優(yōu)化難度較大[2]。而LQCD的矩陣非常規(guī)則，數據的并行切分和矩陣元的構造都很方便直觀，有較大的潛在優(yōu)化空間。因此類似于很多其他差分問題，LQCD的矩陣操作往往寫成模板計算（stencil）的形式，只不過是更復雜的四維九點的模板計算。在具體的實現(xiàn)中，往往采用“內部+外暈”的數據模型，并采用異步通信實現(xiàn)計算與數據傳輸的重疊。矩陣求解算法一般采用Krylov子空間迭代算法，并配合各種預處理算法，如多重網格算法和域分解算法等。這部分的常用算法非常多樣，如LQCD特殊的multi-shift算法、奇偶預處理技術等。

相對來說，張量計算比較清晰簡單，一般可以用愛因斯坦求和記號簡潔地進行描述，在程序中的實現(xiàn)往往是多重循環(huán)配合對稱并行的。但考慮到它實際消耗的計算資源的比例，現(xiàn)有的軟件大多嚴重地低估了這類函數，并沒有給出通用的接口，針對性的優(yōu)化也比較少見。

2 格點量子色動力學計算軟件的現(xiàn)狀

國內外的每個LQCD研究合作組都擁有自己獨特的代碼積累，其中一些較通用的部分往往以開源的方式與其他合作組共享。美國USQCD合作組在美國SciDac經費的支持下開發(fā)的USQCD軟件集是一個相當成熟完善的開源范例，形成了一定程度上的事實標準。這個軟件集包括4層框架，每層有若干特定功能的軟件，它們按照接口協(xié)議相互協(xié)作，可實現(xiàn)復雜且高效的計算功能。

USQCD軟件集的最底層包括負責基本線性代數操作的QLA（QCD linear algebra）、包裝多進程和多線程并行的QMP（lattice QCD message passing）和QMT（QCD support for multi-thread）；第2層包括QDP（SciDAC QFT data parallel library）協(xié)議的各種實現(xiàn)、包裝文件輸入輸出的QIO（QCD input/output applications programmer interface）和c-lime；第3層包括實現(xiàn)稀疏矩陣求解的各類算法；第4層負責整合功能，并提供人機接口，比如Chroma提供了一個集成框架和XML交互協(xié)議，QLua為QDP協(xié)議提供了一個Lua語言的包裝，CPS（columbia physics system）以C++庫包的形式提供了編程接口。

在這個4層框架中，最關鍵的特點是USQCD獨創(chuàng)的QDP協(xié)議。QDP為上層提供了一個方便的編程模型：在這里，規(guī)范場或費米子場等數據可以被自動拆分，并被安排到不同的計算節(jié)點上，所有的操作也都是單程序多數據（SPMD）風格的并行實現(xiàn)，然而這些實現(xiàn)都被包裝在一些C++對象里，上層程序不需要考慮任何有關并行的細節(jié)。考慮到LQCD數據往往都是四維的場，QDP利用四維時空指標切分數據，并提供了全局的平移（shift）操作。QDP協(xié)議有C語言和C++語言的實現(xiàn)，后來又有了采用運行時編譯技術的C++實現(xiàn)，即目前最常用的QDP-JIT。

隨著GPU在高性能計算中的普遍應用，在NVIDIA公司的大力支持下，QUDA逐漸成熟。QUDA在USQCD框架中的定位是第3層，即它充分利用GPU的性能實現(xiàn)了各種稀疏矩陣的高效求解。然而，為了充分利用GPU硬件的特點，QUDA自身實現(xiàn)了大量底層功能，擊穿了下面的兩層協(xié)議，偏離了USQCD框架關于分層解耦的設計初衷。

另外一個流行軟件Grid試圖代替QDP重新構建一個框架。Grid的一個設計特點是面向CPU的向量指令，不僅把四維時空的切分放在不同的進程上，也放在向量指令的內部。這個方法可以實現(xiàn)大量程序的自動向量化，而且對于不同字長的向量指令具有良好的可移植性。后來，隨著GPU的發(fā)展，Grid也提供了對GPU的良好支持，它采用一些宏和預編譯指令精巧地實現(xiàn)了跨平臺運行。

中國格點合作組（China Lattice QCD Collaboration，CLQCD）在近十幾年的科學研究中，積累了很多特定功能的計算代碼，但沒有一個成體系的軟件框架。隨著國產超級計算（以下簡稱超算）的發(fā)展，這個問題在代碼研發(fā)和移植過程中越來越凸顯。

近年來，研究人員基于神威·太湖之光超算從零開始研發(fā)了神威格點原型軟件（SWLQCD），通過優(yōu)化流程細節(jié)、嵌入匯編、手工向量化等細致的工作實現(xiàn)了較理想的運行效率[3-4]。但由于目前沒有國產的軟件框架，只能通過制作QSunway接口與QDP協(xié)議連接，然后把SWLQCD作為一個模塊嵌入USQCD的Chroma軟件內，實現(xiàn)與上層的其他算法協(xié)作運行。

同時，筆者團隊也在“天河三號”原型機上移植了Chroma和Grid，并針對國產處理器的向量指令進行了優(yōu)化[5]；針對曙光的新超算環(huán)境，移植了QUDA，把原本面向CUDA的代碼改寫成面向HIP的代碼[6]。

目前，在國產的3個超算環(huán)境上，都實現(xiàn)了LQCD軟件的研發(fā)或移植。然而，由于軟硬件環(huán)境與國外區(qū)別很大、國產環(huán)境的不完善，開源軟件的移植效率、自主研發(fā)的可持續(xù)性等都是研究者面臨的嚴峻挑戰(zhàn)。

3 編程模型的演化與國產環(huán)境的特點

軟件更新?lián)Q代的動因往往是編程模型的革新，而底層編程模型依賴于硬件的實現(xiàn)。隨著計算機硬件的不斷發(fā)展，對應的編程模型和基礎軟件應運而生，應用層的軟件面臨更替或重構的挑戰(zhàn)。近年來，國產超算的硬件平臺發(fā)展速度非常快，至少有3個主要的架構在并行發(fā)展，而對應的編程模型和基礎軟件尚未跟上步伐，這使得應用層的軟件無所適從，難以做到高效且可移植。在多種異構環(huán)境下，提出新的編程模型，并將其實現(xiàn)為高效且可移植的標準庫，這是國產環(huán)境的軟件生態(tài)發(fā)展的關鍵環(huán)節(jié)。

縱觀LQCD軟件的發(fā)展歷程，這個脈絡是很清晰的。

20世紀末，隨著向量機的退場，Beowulf集群成為大規(guī)模高性能計算和大數據處理的主流平臺。在這樣的集群下，節(jié)點之間通過網絡互聯(lián)，無法直接訪問遠端內存，天然適合多進程和消息傳遞的并行模型。因此，各合作組研發(fā)的LQCD軟件都開啟了MPI（message passing interface）并行化的變革。隨后，片上對稱多核處理器的流行推動了共享內存的線程模型，可以很方便地用OpenMP預編譯指令指導熱點循環(huán)塊的本地并行化。USQCD的QDP協(xié)議就是在這個時候抓住了關鍵的抽象要素——把對稱的數據放在對稱的處理器上，設計了一套針對LQCD計算的SPMD抽象層，從而封裝了單進程、多進程、多線程的各種底層實現(xiàn)。

Grid的設計大約是在Intel公司提出Xeon Phi系列處理器的時候進行的，此時各種不同的單指令多數據流（SIMD）指令對代碼的高效移植已形成了很大的負擔。考慮到SIMD指令的向量維度其實也可被看作對稱多處理器的一個新維度，因此可以把時空點在這個維度上進行切分。這個思路啟發(fā)了Grid的設計與實現(xiàn)。

然而，Intel公司并沒有繼續(xù)發(fā)展融核技術，目前美國最新的超級計算機都是基于CPU+GPGPU架構的。因此，用CUDA編程模型開發(fā)的QUDA逐步走上前臺，成為在GPU上計算LQCD的首選軟件。

隨著C++標準的不斷升級，C++能夠承載的編程模型更加豐富。ROCm/HIP是在C++語法基礎上重建的一套類似CUDA的編程模型和相應的基礎庫，可以適配AMD公司和曙光公司的GPU。未來C++標準中有可能會加入異構計算的支持，目前比較接近的編程模型是SYCL和Kokkos，它們的特點是把匿名函數作為核函數下放給加速設備去運行。這兩個編程模型下的LQCD原型程序已經實現(xiàn)了初步的版本。

然而，我國LQCD軟件的未來發(fā)展路徑需要走向另一個方向，這是由國產軟硬件環(huán)境決定的。

美國的幾個商業(yè)公司分別正在推動CUDA、HIP、SYCL幾個編程模型的互相交叉兼容。因為美國最先進的超算都是基于CPU+GPGPU的硬件模型的，所以編程模型也是類似的。又由于對CUDA的路徑依賴，HIP需要兼容CUDA，SYCL需要兼容CUDA和HIP。它們的編程模型有一個基本的共同之處：都是基于“主-從”模型的，默認的主控代碼運行在CPU端，它可以控制加速設備執(zhí)行熱點函數。

我國的超算有3個并列發(fā)展方向，軟硬件環(huán)境各不相同，編譯器對語法標準的支持也很有限，因此沒必要盲目跟隨SYCL等編程模型。由于沒有路徑依賴的歷史包袱，我國的基礎軟件研發(fā)可以面向更長遠的目標，尋找最適合未來國產超算發(fā)展方向的新編程模型。

國產申威處理器是片上異構的設計，包含若干核組，每個核組包含一個主核與64個從核。軟件環(huán)境提供的是Athread編程接口，這個編程模型是面向硬件特性設計的，因此抽象程度不足。雖然Athread中也用主核調度從核，但研究人員在實際的代碼研發(fā)實踐中發(fā)現(xiàn)，把大部分邏輯留給從核處理是很好的策略，主核只需要被動地做一些數據傳輸和輸入輸出等輔助性的工作。代碼設計的視角和重心在從核一側，這與傳統(tǒng)異構編程的思路有所不同。

日本的超級計算機“富岳”采用的處理器也是片上異構的，每核組也是一個輔助核心和一組計算核心陣列的組合。

實際上，即使是在GPU硬件環(huán)境下，這種不協(xié)調的感覺也已經浮現(xiàn)。由于PCIE（peripheral component interconnect express）的帶寬有限，在多GPU節(jié)點上，NVIDIA公司采用NVLink或NVSwitch來實現(xiàn)GPU之間的直接通信，這種做法繞過了CPU和主內存。另外，統(tǒng)一內存映射和遠程直接數據存取（remote direct memory access，RDMA）也弱化了CPU在數據傳輸中的重要性。因此，未來的代碼會不可避免地向加速設備傾斜，目前流行的以CPU代碼為主線的編程模型并非最優(yōu)之選。

國產的神威、天河、曙光3個系列大規(guī)模超算平臺的共同點是異構。異構意味著計算單元和存儲單元都不是單一、對稱的，它們之間應當有豐富的分工協(xié)作，未必是單純的“主-從”模型。筆者期望，下一代編程模型應該更細致地刻畫不同的處理器和內存資源，并定義出通用的控制權轉移和數據傳輸的規(guī)則和相關接口。

LQCD計算的性能瓶頸經常體現(xiàn)在數據流動上。代碼優(yōu)化的重點經常是節(jié)點間的MPI傳輸的時機、CPU緩存的合理復用、GPU顯存的充分安排、申威從核直接內存訪問（direct memory access，DMA）的錯峰傳輸等。然而，目前的高性能計算的編程模型很少有面向數據流動和異構內存空間的支持，絕大多數細節(jié)需要手動安排。申威和NVIDIA公司都為各自的硬件實現(xiàn)了OpenAcc編程模型，這對于初學者而言入門方便，但也阻擋了優(yōu)化的可能性，無法支撐優(yōu)秀的應用。筆者期望，下一代的編程模型應當在數據流動方面做出更合理的抽象，使得這方面的優(yōu)化可以被自動或半自動地完成。

4 應用軟件層與基礎軟件層的新編程模型

把前述兩個對下一代編程模型的期望進行綜合考慮，筆者發(fā)現(xiàn)符合要求的編程模型難以用庫包或語法設計實現(xiàn)，它很可能更像一個大數據處理引擎，只不過它的粒度比Spark等傳統(tǒng)引擎要細得多，要直接深入硬件細節(jié)中，才能保證極限的計算效率。為了配合一個引擎風格的編程模型，應用層的程序需要被拆分成很多獨立的算子，這些算子根據數據依賴關系構成了有向無環(huán)圖，被引擎依序分配給不同的資源進行計算，這種做法與傳統(tǒng)的大數據處理引擎相同。

由于算子的粒度較小，算子切換的代價需要被極限優(yōu)化，這在不同硬件環(huán)境上可能需要完全不同的優(yōu)化技術。在國產申威架構下，從核有能力獨立處理復雜的任務切換，但多從核互相配合進行異步操作時的靈活性不足。這是因為從核是單線程的，沒有天然的并發(fā)支持。為此，筆者團隊做了一個試驗性的代碼OSP，它是在程序塊粒度上的超輕量級協(xié)程庫，能夠支持多個偽多線程并發(fā)，并提供了私有變量支持。雖然筆者團隊采用了一些非常特殊的語法技巧，但全部代碼完全符合C89標準語法，可移植性非常高。筆者團隊利用OSP在申威平臺上為從核DMA做了異步包裝，初步測試成功。與此同時，筆者團隊還設計了一套被稱作虛擬機生成器（CVM）的優(yōu)化框架，采用遺傳算法為64個申威從核尋找最優(yōu)的靜態(tài)任務調度方案，相比人工方案，優(yōu)化了13%的效率。這兩項工作分別成功驗證了動態(tài)調度和靜態(tài)調度的可行性，為今后在各類國產平臺上建立輕量級的任務調度引擎奠定了基礎。

與粗粒度的大數據處理引擎不同，貼近硬件的細粒度引擎需要細致地考慮各種不同內存之間的數據傳輸，包括主內存（可能有核心綁定）、多級緩存、GPU的多種不同訪問模式的顯存、申威從核便箋存儲器（LDM）等。目前不同軟硬件環(huán)境的接口各不相同，沒有統(tǒng)一的公共編程模型，這是需要努力創(chuàng)新的地方。

一般來說，通信有兩種效果：數據傳輸和控制流同步。從形式上看，MPI和申威寄存器通信等基于消息傳遞模型的方式是顯式的數據傳輸、隱式的同步；而線程模型的共享內存方式是顯式的同步、隱式的數據傳輸。目前看來，后者似乎更有希望成為未來高性能計算的主流，MPI-3版本也增加了共享內存的接口。考慮到訪問局部性對效率至關重要，完全的統(tǒng)一編址并不符合人們的期待。在目前已有的訪問模型中，遠程內存訪問（remote memory access，RMA）可能是比較適用且高效的。

那么，新的編程模型需要提供對不同內存的統(tǒng)一描述，對遠程數據的讀寫提供統(tǒng)一的接口。異步的遠程讀寫往往會有一些標志用于判斷傳輸完成情況，這就給任務調度引擎提供了一個重要的信息：它可以根據當前數據的完備情況安排算子的運行。

在這樣的編程模型下，應用層的軟件編寫應當分為兩部分：一部分是各類算子的代碼實現(xiàn)，另一部分是描述性地列出數據依賴關系和傳輸的拓撲關系。這種編程模型與目前常見的模型迥然不同，它是面向數據的，代碼是連接數據的橋梁，是零散的。

這樣的編程模型有利于極限編譯優(yōu)化和高效運行，但對于大部分程序員來說并不友好。因此，還需要另外一個面向應用層的編程模型，把底層模型進一步抽象化，讓它符合用戶的思維方式。

回顧LQCD的計算需求：大部分計算熱點函數可以描述為張量表達式。求解大規(guī)模稀疏矩陣時，此稀疏矩陣作用在向量上的操作也可以寫成張量表達式的形式——在物理理論里就是如此表述的。從計算機程序的角度看，稠密的張量就是多維數組，稀疏的張量往往是簡單的枚舉函數，用這兩種簡單的元件可以構筑出LQCD絕大多數的計算過程。可以進一步猜想，張量運算形式的描述能力可能覆蓋了各領域的科學計算中的很大一部分，做好張量計算的封裝和優(yōu)化可能會為科學計算提供強有力的支持。

在LQCD的計算實踐中，因處理數據的需求，筆者團隊設計并實現(xiàn)了QScheme語言。它在Scheme語言的基礎上，增加了xdata數據格式和對應的xfile文件格式。這兩種格式本質上就是“帶有元數據的多維數組”，這個元數據包括數組的維數、每個維度的名稱、每個維度的指標列表。筆者團隊為xdata實現(xiàn)了豐富的操作功能，這些操作的基礎是xdata四則運算的規(guī)則：

● 運算結果的維度集合是各參數的維度集合的并集；

● 運算結果的每個維度的指標集合是含有這個維度的參數的對應指標集合的交集；

● 程序實現(xiàn)保證對各參數維度的對應指標分別進行計算, 與它們在內存中的存儲順序無關；

● 普通數值被視作包含0個維度的 xdata。

這個規(guī)則可以被視為擴展的張量運算規(guī)則（與數學規(guī)則相比，多了一個“自動丟棄無法匹配的指標”的例外規(guī)則），可以把那些平時不被視為張量分量的維度納入張量處理的框架中，進一步提高了模型的適用性。

筆者團隊使用QScheme進行了長時間的實際科研工作，發(fā)現(xiàn)xdata的表達能力強、優(yōu)化方便，非常適合科學計算和數據處理。因此，前述的底層編程模型如果能夠進一步被抽象為類似xdata的張量計算模型，會非常適合科研人員使用和進行二次開發(fā)。

張量表達式的形式簡單、規(guī)則明確，可以進行靜態(tài)分析，進而能夠用程序估算出表達式的計算資源需求，包括浮點操作計數、內存占用等預算信息。可以基于項或指標拆分表達式，利用子表達式的資源預算找到最優(yōu)的拆分方案，然后用元編程技術把子表達式輸出為算子的計算程序，把數據拆分情況整理成數據描述信息，最后把這些信息傳遞給下層的調度引擎。這樣，兩層計算模型之間就可以用一套靜態(tài)分析算法和代碼生成算法連接起來。

至此，筆者團隊期待的高性能大數據處理的基礎軟件已逐漸成形：它包含上下兩層編程模型，下層實現(xiàn)高效跨平臺的遠程內存訪問接口和算子調度引擎；上層是在此基礎上構建的擴展張量計算模型，為研究者提供張量表達式形式的代碼接口。

5 結束語

本文回顧了LQCD的計算特點和軟件沿革，基于國產超算的軟硬件環(huán)境和發(fā)展方向提出了面向LQCD和其他科學計算的上下兩層編程模型。筆者團隊研發(fā)了試驗性的軟件，用原型測試和實際使用的經驗為這個設想中的編程模型提供了可行性的支撐。這個編程模型的實現(xiàn)既是國產環(huán)境基礎軟件的一項挑戰(zhàn)，也是繁榮軟件生態(tài)的一個可能的契機。