中文題目：Mille-feuille：GPU上的分塊混合精度單內核共軛梯度解法器

論文題目：Mille-feuille: A Tile-Grained Mixed Precision Single-Kernel Conjugate Gradient Solver on GPUs

錄用期刊/會議：37th International Conference for High Performance Computing, Networking, Storage, and Analysis (CCF A)

錄用/見刊時間：2024-6-15（錄用時間）

作者列表：

1）楊德闖 中國石油大學（北京）人工智能學院 計算機技術 碩21

2）趙雨軒 中國石油大學（北京）人工智能學院 計算機科學與技術 碩22

3）牛一多 中國石油大學（北京）人工智能學院 計算機科學與技術 本21

4）賈偉樂 中國科學院計算技術研究所 研究員

5）邵   恩 中國科學院計算技術研究所 高級工程師

6）劉偉峰 中國石油大學（北京）人工智能學院 計算機系教師

6）譚光明 中國科學院計算技術研究所 研究員

7）金   洲 中國石油大學（北京）人工智能學院 計算機系教師

摘要:

共軛梯度法（CG）和雙共軛梯度穩定法（BiCGSTAB）是用于求解稀疏線性系統的有效方法。本文提出了一種新的求解器——Mille-feuille，用于加速GPU上的CG和BiCGSTAB?；贜VIDIA A100和AMD MI210的實驗結果表明，Mille-feuille求解器在CG中相比基準實現（包括廠商支持的cuSPARSE/hipSPARSE以及兩個最先進的庫PETSc和Ginkgo）取得平均3.03倍/2.68倍，5.37倍，4.36倍（最高可達8.77倍/7.14倍，16.54倍，15.69倍）的加速比；在BiCGSTAB中，平均加速比為2.65倍/2.32倍，3.57倍，3.78倍（最高可達7.51倍/6.63倍，16.64倍，11.73倍）；在預條件CG（PCG）中，平均加速比為3.82倍/3.47倍（最高可達40.38倍/47.75倍）；在預條件BiCGSTAB（PBiCGSTAB）中，平均加速比為1.79倍/1.63倍（最高可達45.63倍/44.34倍）。

背景與動機:

在迭代法解法器中，共軛梯度方法和穩定雙共軛梯度方法分別因在處理對稱正定、非對稱或不正定線性系統中的有效性而突出。目前的迭代法解法器往往忽略了一些對性能至關重要的因素，如稀疏矩陣中非零元素數值精度的分布、內核間的同步開銷以及解向量x的部分收斂等。本文擬利用稀疏矩陣的數值特征、硬件平臺特性以及算法本身的實現原理對共軛梯度解法器進行優化。

設計與實現:

一、稀疏存儲格式

我們設計了一種細粒度的兩級分塊稀疏矩陣存儲格式，高級存儲以COO的結構捕獲塊間的信息用于確保SpMV計算過程的負載均衡，并且根據稀疏矩陣中非零元素初始值的范圍使用四種不同精度進行存儲，對于低級存儲，我們使用 CSR 方式來記錄塊內信息，使用額外的數組記錄塊內非空行的信息，用于避免在 SpMV 操作期間遍歷塊中的空行，從而進一步提高 SpMV的性能。

圖1. 兩級分塊的存儲結構

二、單內核共軛梯度解法器

接下來，為了減少GPU上不同內核之間的同步開銷，我們利用原子操作使整個迭代法求解過程在單個GPU內核內運行。為了實現同步，我們構建多個依賴數組來定義數據和操作之間的依賴關系，并允許原子操作調度warp以執行不同操作的任務。我們根據依賴關系將CG算法分為四個部分，內核啟動之前，矩陣A的非空塊以負載均衡的方式分配給每個warp并且被加載到共享內存中一次，并在迭代過程中被重復利用，這樣可以提升程序的訪存效率。

圖2. 單內核共軛梯度解法器

三、部分收斂感知混合精度策略

最后，為了在迭代過程中利用解向量x 中已經收斂的元素來優化SpMV的性能，我們根據收斂閾值ε設定了四個范圍并實現了一種部分收斂感知混合精度策略，運行時在單個內核內實現塊粒度的片上動態精度轉換。我們的精度轉換僅在共享內存中進行一次，減少了訪問全局內存或執行精度轉換的高昂開銷。

圖3. 部分收斂感知混合精度策略

實驗結果及分析:

一、Mille-feuille與基準實現的對比（cuSPARSE/hipSPARSE）

在 CG 算法上，與基準實現相比，我們的算法的平均加速比為 3.03 倍和 2.68 倍（最高分別為 8.77 倍和 7.14 倍）。在 BiCGSTAB 算法上，與基準實現相比，我們的算法的平均加速比為 2.65 倍和 2.32 倍（最高分別為 7.51 倍和 6.63倍）。

在PCG算法上，與基準實現相比，我們的算法的平均加速比為 3.82 倍和 3.47 倍（最高分別為 40.38 倍和 47.75 倍）。在PBiCGSTAB 算法上，與基準實現相比，我們的算法的平均加速比為 1.79 倍和 1.63 倍（最高分別為45.63倍和44.34倍）。

圖4. 與cuSPARSE和hipSPARSE實現的CG和BiCGSTAB基準算法的性能比較

圖5. 與cuSPARSE和hipSPARSE實現的PCG和PBiCGSTAB基準算法的性能比較

二、Mille-feuille與PETSc和Ginkgo對比

對于 CG 方法，我們的算法相較于 PETSc 和 Ginkgo分別實現了 5.37 倍和 4.36 倍的幾何平均速度提升（最高分別可達 16.54 倍和 15.69 倍）。對于 BiCGSTAB 方法，我們的算法相較于 PETSc 和 Ginkgo 分別實現了3.57倍和 3.78 倍的幾何平均速度提升（最高分別可達 16.64 倍和 11.73 倍）。

圖 6：與PETSc和Ginkgo的性能比較

通訊作者簡介:

金洲，副教授.中國石油大學（北京）計算機系副教授，入選北京市科協青年人才托舉工程、校青年拔尖人才。主要從事集成電路設計自動化（EDA）、面向科學計算的DSA軟硬件協同設計等方面的研究工作。主持并參與國家自然科學基金青年項目、重點項目，科技部重點研發微納電子專項、高性能計算專項青年科學家項目，國家重點實驗室開放課題、企業橫向課題等。在DAC、TCAD、TODAES、SC、PPoPP、IPDPS、TCAS-II、ASP-DAC等重要國際會議和期刊上發表60余篇高水平學術論文。獲EDA2青年科技獎、SC23最佳論文獎、ISEDA23榮譽論文獎、IEEJ九州支部長獎等。

聯系方式：jinzhou@cup.edu.cn