基于申威1600的3級BLAS GEMM函數優(yōu)化①

劉 昊, 劉芳芳, 張 鵬, 楊 超, 蔣麗娟(中國科學院 軟件研究所, 北京 0090)(中國科學院大學, 北京 0090)

基于申威1600的3級BLAS GEMM函數優(yōu)化①

劉 昊1,2, 劉芳芳1, 張 鵬1,2, 楊 超1, 蔣麗娟1,21(中國科學院 軟件研究所, 北京 100190)2(中國科學院大學, 北京 100190)

BLAS是當前科學計算領域重要的底層支持數學庫之一, 其中的3級BLAS函數應用最為廣泛. 本文基于國產申威1600平臺, 提出了一種基礎線性代數庫BLAS的三級函數通用矩陣乘GEMM的高性能實現方法. 在單核上, 使用乘加指令、循環(huán)展開、軟件流水線指令重排、SIMD向量化運算、寄存器分塊技術等與平臺架構相關的技術手段, 實現匯編級手工優(yōu)化; 在多核上, 提出了適用于該平臺的多線程加速方案. 實驗結果顯示, 在單核串行性能測試中, 與知名開源數學庫GotoBLAS相比, 我們實現了平均4.72倍的加速效果; 在多核并行擴展測試中, 4線程版的性能則平均達到了單線程版性能的3.02倍.

申威1600; 三級BLAS; GEMM; 高性能計算; 多核

1 引言

1.1 背景介紹

BLAS(Basic Linear Algebra Subprograms)是一個線性代數核心子程序的集合,主要包括向量和矩陣的基本操作. 它是最基本和最重要的數學庫之一, 廣泛應用于科學工程計算. 目前世界上有關矩陣運算的軟件幾乎都調用BLAS數學庫; 重要的稠密線性代數算法軟件包(如EISPACK、LINPACK、LAPACK和ScaLAPACK等)的底層都是以BLAS為支撐. BLAS目前已成為線性代數領域的一個標準API庫.

BLAS分成三個等級:

1級: 進行向量-向量操作;

2級: 進行向量-矩陣操作;

3級: 進行矩陣-矩陣操作.

本文主要研究BLAS 3級中的通用矩陣相乘GEMM函數, 如公式(1)所示, 其中A, B, C表示矩陣,α，β為標量因子.

針對特定的體系結構進行BLAS庫的優(yōu)化, 特別是GEMM函數的優(yōu)化工作, 國內外已有相當多的研究成果, 例如Goto等人所提出的矩陣乘法實現算法, 通過對矩陣A和B進行劃分, 矩陣乘法細分為一組panel-panel的乘法操作(GEBP), 分析矩陣計算三重循環(huán)順序與cache緩存之間的關系模型并提出最優(yōu)算法[1][2]. 基于該模型的開源高性能BLAS數學庫有Texas大學超級計算中心高性能計算組開發(fā)的GotoBLAS[3].近幾年來, 針對BLAS的自動調優(yōu)技術日漸成熟, ATLAS數學庫可以針對通用CPU平臺進行自動調優(yōu),通過自動調整矩陣分塊、監(jiān)測cache利用率、調整寄存器使用策略等方式, 生成最合適平臺架構的代碼[4].

近年來隨著CPU+GPU異構并行計算的興起, 基礎線性代數庫在這些異構平臺上的高性能實現也出現大量優(yōu)秀科研成果[5,6]. 與本文所述的基于稠密矩陣的BLAS優(yōu)化工作不同, 很多基于STENCIL計算而出現的稀疏矩陣若以稠密方式進行計算則明顯效率低下,面向這些稀疏矩陣實現高性能的BLAS庫是近年來該領域的另一個重要方向[7].

神威藍光國產超級計算機是由國家并行計算機工程研究中心開發(fā)研制的首臺全部采用國產CPU構建的千萬億次超級計算機系統(tǒng), 處理器采用自主設計生產的申威1600 CPU, 于2011年10月份在國家超級計算濟南中心投入使用, 其布局按照MPP萬萬億次架構設計, 主頻975MHz-1.1GHz, 單CPU 16核心[8], CPU總數為8704, CPU核心總數接近14萬, 峰值計算速度達到每秒1100萬億次浮點計算, 持續(xù)計算能力為738萬億次.

隨著高性能計算需求的日益增大, 越來越多的應用開始部署在以申威處理器為代表的國產高性能計算平臺上. 這些應用中, 無論是在傳統(tǒng)的科學計算、氣象預報、金融統(tǒng)計等領域, 亦或是現今新興的機器學習等, 均嚴重依賴底層BLAS庫的實現以獲得較高的計算性能. 而現有的開源BLAS庫, 均無法有效的利用國產平臺的硬件特性, 實際性能較低, 迫切需要針對國產平臺進行優(yōu)化實現的BLAS庫出現. 因此, 結合神威藍光等當前國產多核處理器體系結構特征, 研究BLAS數學庫, 特別是3級BLAS函數的高性能實現方法具有重要意義.

1.2 本文組織結構

本文主要針對國家超級計算濟南中心的申威1600平臺進行BLAS數學庫三級函數的優(yōu)化研究工作,以通用矩陣相乘函數GEMM為例進行表述. 其中第一章是引言, 介紹平臺信息與BLAS庫優(yōu)化現狀以及本文工作的意義; 第二章以GEMM為例, 介紹三級BLAS函數特點; 第三章介紹基于單核的GEMM優(yōu)化設計; 第四章介紹基于多核的GEMM優(yōu)化設計; 第五章為實驗, 與開源的GotoBLAS數學庫進行性能比較;最后是總結與展望.

2 三級BLAS GEMM函數概述

一般來說, 1級BLAS函數計算復雜度為O(N), 2級BLAS函數計算復雜度為O(N2), 3級BLAS函數計算復雜度為O(N3), 評價一套BLAS數學庫性能高低, 3級BLAS的性能是最重要的指標, 而通用矩陣乘法函數GEMM作為3級BLAS最常用的函數, 其性能指標往往也最被關注.

GEMM實現公式1所述的矩陣相乘計算, 其本質是K-N-M的三重循環(huán)計算, 計算復雜度為O(N3). 考慮其訪存特性, GEMM函數讀取A、B、C矩陣各一次,寫回C矩陣一次, 復雜度可表述為O(2MN+MK+KN),整體為O(N2)復雜度, 在矩陣的K、N、M三維足夠大的情況下, 計算復雜度遠大于訪存復雜度, 將訪存開銷隱藏在計算過程中是完全可行的, 理論上GEMM函數性能應可以接近平臺性能峰值.

目前的通用編譯器往往只能進行普適性的優(yōu)化,例如循環(huán)展開, 利用硬件流水線進行加速, 而不能針對GEMM函數特性, 分析數據依賴情況, 實現最優(yōu)化的計算與訪存延遲互相掩蓋. 例如著名的NETLIB BLAS, 一種通用的, 使用FORTRAN語言實現的BLAS庫(一般作為正確性與性能的基準測試程序, http://www.netlib.org/blas/). 該BLAS庫中的GEMM使用最樸素的三重循環(huán)實現, 可以認為除了編譯器優(yōu)化外不存在其它任何形式優(yōu)化. 該版本BLAS在本文所述的申威1600平臺上性能為280 MFLOPS, 僅僅相當于平臺性能峰值的3.5%, 遠不能滿足應用對GEMM的性能要求[9], 故需要設計高效的分塊算法, 結合核心段的需要手工匯編來實現.

3 基于單核的串行優(yōu)化設計

針對申威1600平臺單核, 基于其平臺的擴展ALPHA指令集, 本節(jié)提出了GEMM函數的高性能優(yōu)化方案. 該方案基于本文第2章所述的GEMM函數的計算密集型特征, 充分考慮申威1600平臺特性, 最大程度發(fā)揮該平臺計算能力, 實現該平臺上的高性能GEMM函數, 合理利用平臺的寄存器等硬件資源, 實現計算與訪存延遲互相掩蓋.

3.1 乘加指令

GEMM函數實現的計算如公式1所示, 對C矩陣元素進行乘加操作. GotoBLAS針對ALPHA架構的核心函數的匯編實現中沒有實現乘加操作, 乘法與加法分別進行. 在寄存器分塊為4*4的情況下, 讀取A、B矩陣各4份數據, 計算16份C矩陣的部分解, 計算訪存比達到4比1, 在內層循環(huán)充分循環(huán)展開后, 平均每4條計算指令中插入一條通信指令. 高達4比1 的計算通信會增加寄存器的消耗, 不利于軟件流水線的排布.

申威1600平臺提供SIMD乘加指令, 每次計算指令可以實現加法和乘法在一個指令周期完成. 在同樣4*4的寄存器分塊情況下, 計算通信比降低為2比1[10],計算指令總數減少一半, 理論上可達到2倍的性能加速, 內層循環(huán)中的中間變量數減少, 對寄存器的消耗降低, 利于軟件流水線的排布.

3.2 256位的SIMD向量運算

一般而言, 在不考慮編譯優(yōu)化的情況下, 我們可以使用intrinsic函數進行手工向量化, 基于C代碼就可以取得不錯的性能提升. 但是對于性能要求較高的任務, 比如本文討論的GEMM函數, intrinsic函數并不能勝任, 因為intrinsic函數不能控制指令的調度和寄存器分配, 所以我們需要在核心函數中編寫向量化的匯編代碼[9].

申威1600平臺提供了256位的SIMD向量擴展支持. 以實數雙精度為例, 每個元素為8字節(jié)、64位, 進行向量化擴展后, 配合長度為256位的浮點向量化寄存器, 每次順序讀取4個A矩陣元素, 同時對4個C矩陣元素進行運算, 理論上性能可實現4倍加速.

3.3 循環(huán)展開與軟件流水線

循環(huán)展開是編譯器經常使用的優(yōu)化策略之一, 循環(huán)展開后, 最內層循環(huán)的單次循環(huán)內指令數增加, 從硬件角度看, 更多的指令可以用于亂序發(fā)射與亂序執(zhí)行, 利于硬件流水線發(fā)揮作用, CPU保留站等硬件資源也可更充分得以利用; 從軟件角度看, 更多的指令可以用于指令重排, 有助于避免“訪存-計算”的數據依賴關系. 一般而言, 訪問主存往往有較高的延遲, 要求多次的循環(huán)展開, 將訪存操作的延遲隱藏于運算操作中.

在實際的核心函數實現中, 我們仍然采用了經典的N-M-K的計算順序, 內層循環(huán)的部分代碼如下, 其中VMAX為乘加指令:

圖1 核心匯編代碼舉例

在核心函數設計過程中, 以下幾點比較重要: (1)硬件提供兩條流水線, 無數據依賴的計算與訪存可同時發(fā)射, 則指令排布應盡量保證一條計算指令和一條取數指令共同出現, 使得兩條流水線都占滿; (2)最內層循環(huán)需展開足夠的次數, 掩蓋掉所有的本地訪存指令, 考慮實際申威1600平臺數據從cache到寄存器的延遲, 最內層的K層循環(huán)展開兩次即可.

3.4 寄存器分塊

寄存器分塊的目的是合理安排核心函數中矩陣分塊的大小, 在不超過硬件限制的情況下盡可能充分利用寄存器等硬件資源, 實現性能優(yōu)化.

寄存器的數量決定了寄存器分塊的大小. 申威1600是擴展的ALPHA架構平臺, 每個CPU擁有F0到F31, 共32個邏輯浮點寄存器, 本文選擇4*4的寄存器分塊. 在此分塊下, 每輪使用4個浮點向量寄存器, 取4個256位向量長度的A矩陣元素, 共16個A矩陣元素; 使用4個浮點向量寄存器, 每個寄存器取1個B陣元素并擴展成4份, 則4個寄存器處理4個B矩陣元素, 計算16個256位向量長度的C矩陣元素,即64個C矩陣雙精元素的部分解, 由此共使用了24個浮點向量寄存器. 此外, F31寄存器固定存儲0值,剩余7個邏輯寄存器存儲中間變量或用作循環(huán)變量,寄存器資源使用合理且高效.

3.5 數據預取與數據重排

申威1600提供特定的預取指令fetchd, 支持將內存中的數據預取至cache中, 每次預取一個cache行的數據量. 若不使用預取指令, 每當一個cache行的數據使用完畢, 硬件都需要等待數據從內存加載至cache中,造成流水線停頓, 且這種停頓以百拍來衡量, 對于追求高性能的GEMM來說是無法容忍的. 在實測中, 若矩陣A拷貝過程不使用預取指令, 總體性能下降14.1%,若矩陣B拷貝過程不用預取指令, 總體性能下降43.6%.

使用fetchd預取指令的位置及預取跨步的大小需謹慎設計, 若跨步太小, 會造成數據還未完成預取就需要使用, 造成流水線停頓; 若跨步太大, 會造成預取數據完成后長時間沒有使用, 再次被替換出去. 申威1600 CPU使用的cache替換順序為FIFO, 即先進先出的策略, 這對GEMM函數的設計而言是有利的, cache數據的替換是可控的, 因此只要測試出合適的跨步, 性能可以保證穩(wěn)定.

數據的預取僅考慮到指令位置與跨步大小是不夠的. 按照Goto等人的分析[1,11], GEMM核心函數三重循環(huán)的最優(yōu)順序應是N-M-K, 且最內層循環(huán), 即K層循環(huán)應保證足夠高的性能, 盡量減少延遲, 最大程度利用好cache和寄存器等硬件資源, 否則函數整體性能無從談起[12].

我們以列主元矩陣為例進行說明, 若矩陣數據不進行任何處理直接按照N-M-K的順序進行計算, 使用前文所述的4*4寄存器分塊, 不考慮SIMD向量化影響, 且假設分塊矩陣的M維足夠大, 則每次cache行只有前4個元素能夠使用, 沿K維進行的下一次計算需要的數據一定不在該cache行中, cache資源被極大浪費, 且極易因cache不命中造成流水線停頓.

為此, 需考慮將A矩陣和B矩陣沿K維方向進行重排, 每次cache行所取數據都被完全利用, 重排效果如圖2所示.

圖2 A矩陣和B矩陣重排示意圖

4 基于多核的優(yōu)化設計

考慮計算任務的劃分, 各線程對矩陣進行等份劃分, 計算各自對應的C矩陣部分解即可, 理論上講, 線程之間沒有數據依賴關系, 可以說GEMM函數具有天然的多線程負載均衡特性. 對于GEMM函數而言, 進行多核并行化優(yōu)化實施并不復雜, 每個計算核心運行一個線程, 每個線程之上的優(yōu)化設計仍采用上一章所述的基于單核的優(yōu)化設計方案即可.

直觀來看, 任務劃分可以按照A矩陣的M維進行劃分或按照B矩陣的N維進行劃分, 本文使用了行方向, 即沿著A矩陣的M維進行劃分的方法, 4線程并行化的數據劃分, 如圖3所示, 每個線程計算對應的C矩陣部分解.

如圖3所示, 4個線程各自擁有私有的部分A矩陣,線程間共享整個B矩陣, 計算對應的C矩陣的部分解.由于申威1600架構支持多線程共享內存, 理論上4線程對一個內存地址具有相同的訪存延遲, 不會發(fā)生某個線程需要等待其他線程訪存的情況.

正如前文所述, GEMM具有天然的負載均衡特性,在矩陣規(guī)模合理, 不考慮特殊規(guī)?；蜻吔绲那闆r下,可以做到所有線程計算規(guī)模完全或近似相等. 然而,由于每個線程內部的計算仍然按照單核進行, 仍需采用寄存器分塊、軟件流水線等技術, 每個線程需各自進行A、B矩陣的數據重排, 我們對A、B矩陣的重排情況分別進行討論.

圖3 線程間數據劃分

首先考察A矩陣, 由于每個線程私有部分A矩陣元素, 線程內部進行的矩陣重排不會影響到其他線程,多個線程可以同時進行各自A矩陣的數據重排.

接下來考察B矩陣, 我們已知所有線程共享B矩陣的所有元素, 一種簡單有效的方案是所有線程等待1號線程完成B矩陣的數據重排, 然后再并發(fā)進行計算. 但是這種方案破壞了GEMM本身具有的負載均衡特性, 在1號線程進行數據重排的過程中, 其他線程將進行等待; 另一方面講, B矩陣的劃分工作從直觀上講是完全可以并發(fā)執(zhí)行的. 因此, 我們實際也需要將B矩陣劃分于各個線程之中, 如圖4所示, 將B矩陣按列劃分于每個線程之中.

圖4 B矩陣劃分

考慮到B矩陣需要共享, 我們設計了一種二維數組的線程間共享的數據結構working[i], 表示第j個線程為第i個線程準備的重排后的B矩陣元素, working[i]中的每個元素為一個指針, 指向對應的重排后的B矩陣元素頭指針. 換言之, 重排后的B矩陣數據被各個線程私有, 但這些數據指針卻是共享的, 基于申威1600的平臺多線程共享內存架構, 各線程可以相同的延遲共享這些數組.

該共享的數據結構內的元素為指針, 可用于線程間同步而無需另外設計線程間的鎖機制, 當working[i]為空時, 表示此時的對應B矩陣還未重排準備就緒,不能用于計算, 該線程應該等待; 反之, 表示該B矩陣已重排準備完畢, 第i個線程可以獲取第j個線程的B元素進行計算, 計算完畢后working[i]置為空, 等待下一輪的第j個B矩陣重排完畢. 同時, 只有working[all]都為空時, 第j個線程的B矩陣才算使用完畢, 沒有其他線程需要這部分數據可以釋放并進行下一輪的B矩陣數據重排.

5 性能測試與評價

申威1600平臺作為國產高性能多核平臺, 性能指標優(yōu)異, 已被國家超級計算濟南中心等國內超算中心使用, 該平臺即為本文實驗的測試平臺. 該平臺使用了擴展的ALPHA架構指令集, 每個CPU支持乘加指令與256位向量化運算, 支持計算指令與訪存指令同步發(fā)射, 具體硬件參數如表1所示.

表1 申威1600主要參數

表2統(tǒng)計了統(tǒng)計基于單核的單線程方案, 數據規(guī)模分別為1024、2048、4096、8192, 申威1600 GEMM函數性能和GotoBLAS GEMM性能, 如圖5所示. 實驗涵蓋了如上四種規(guī)模的實數單精、實數雙精、復數單精、復數雙精共計16組測試用例, 平均加速比為4.72, 最高加速比為5.61. 為說明基于申威多核的4線程設計方案加速效果, 表3統(tǒng)計了基于多核的4線程方案, 包含了數據規(guī)模與數據類型的同上的16組測試用例, 與申威平臺單核的單線程方案相較, 平均加速比為3.02, 最高加速比為3.18. 對比圖如圖6所示.

表2 申威單核與GotoBLAS單核性能對比

表3 申威多核4線程與申威單核性能對比

6 總結與展望

本文首先說明了三級BLAS GEMM函數的計算密集型的特征, 接下來介紹了國產申威1600多核平臺的特征. 基于該平臺擴展的ALPHA架構, 提出了基于該平臺的高性能BLAS三級函數實現方案. 本方案首先分析了基于單核的單線程優(yōu)化方案, 提出了矩陣數據重排、軟件流水線、寄存器分塊等基于該平臺的優(yōu)化技術;接著分析了基于多核的多線程優(yōu)化方案, 提出了矩陣劃分的方式與多線程同步機制. 在性能測試中, 基于申威1600平臺優(yōu)化過的GEMM函數性能, 在單核情況下對比GotoBLAS的平均加速比為4.72, 多核4線程方案的平均加速比為3.02, 達到了預期的效果.

圖5 申威單核與GotoBLAS對比

圖6 申威4線程與單核對比

申威1600作為新興的國產高性能計算平臺, 我們期待著該平臺系列的后續(xù)機型, 針對該系列平臺未來將可能出現的新特性, 這對這些特性又會對高性能的BLAS數學庫提出怎樣的要求, 我們拭目以待.

1 Goto K, Geijn RA. Anatomy of high-performance matrix multiplication. ACM Trans. on Mathematical Software (TOMS), 2008, 34(3): 12.

2 Goto K, Van De Geijn R. High-performance implementation of the level-3 BLAS. ACM Trans. on Mathematical Software (TOMS), 2008, 35(1): 4.

3 蔣孟奇,張云泉,宋剛,等.GOTOBLAS一般矩陣乘法高效實現機制的研究.計算機工程,2008,34(7):84–86,103.

4 Whaley RC, Petitet A, Dongarra JJ. Automated empirical optimizations of software and the ATLAS project. Parallel Computing, 2001, 27(1-2): 3–35.

5 張帥,李濤,王藝峰,等.細粒度任務并行GPU通用矩陣乘.計算機工程與科學,2015,37(5):847–856.

6 Kurzak J．Preliminary results of autotuning GEMM kernels for the NVIDIA Kepler architecture-GeForce GTX 680 [z3.I, APACK Working Note 267,2014

7 李佳佳,張秀霞,譚光明,等.選擇稀疏矩陣乘法最優(yōu)存儲格式的研究.計算機研究與發(fā)展,2014,51(4):882–894.

8 申威1600處理器的細枝末節(jié).黑龍江科技信息,2011, (34):I0004–I0004.

9 張先軼.若干線性解法器多核和眾核性能優(yōu)化關鍵技術研究[學位論文].北京:中國科學院大學,2014.

10 李毅,何頌頌,李愷等.多核龍芯3A上二級BLAS庫的優(yōu)化.計算機系統(tǒng)應用,2011,20(1):163–167.

11 張先軼,王茜,張云泉,等.OpenBLAS:龍芯3A CPU的高性能BLAS庫.2011年全國高性能計算學術年會(HPC china2011)論文集,2011.

12 陳少虎.BLAS庫在龍芯3A上的實現與優(yōu)化[學位論文].北京:中國科學院研究生院,2011.

Optimization of BLAS Level 3 Functions on SW1600

LIU Hao1,2, LIU Fang-Fang1, ZHANG Peng1,2, YANG Chao1, JIANG Li-Juan1,212
(Institute of Software, Chinese Academy of Sciences, Beijing 100190, China ) (University of Chinese Academy of Sciences, Beijing 100190, China)

BLAS is one of the most important basic underlying math library for scientific computing, in which the level 3 BLAS functions are most widely used. In this paper, we provide a high-performance method to implement Level 3 BLAS functions based on domestic Sunway 1600 platform. To make it clear, we take GEMM as an example. For the implementation on single-core, we apply many tuning techniques related to the specific platform, such as multiply-add instructions, loop unrolling, software pipelining and instruction rearrangement, SIMD operations, and register blocking to push up the performance. For the multi-core implementation, we propose an efficient multi-threaded method. Compared with GotoBLAS, one of the famous open-source BLAS, the experiments show that our serial single-threaded method achieves a speedup of 4.72. What’s more, the average speedup of 4-threaded execution towards the single-threaded one can also reach 3.02.

Sunway 1600; level 3 BLAS; GEMM; HPC; multi-core

國家自然科學基金(91530103,91530323)

2016-03-16;收到修改稿時間:2016-04-27

10.15888/j.cnki.csa.005456