中文題目：AmgT：Tensor Core加速的AMG解法器

論文題目：AmgT: Algebraic Multigrid Solver on Tensor Cores

錄用期刊/會議：37th ACM/IEEE International Conference for High Performance Computing, Networking, Storage, and Analysis. (CCF A) Best Paper Finalist. 最佳論文提名

原文DOI：10.1109/SC41406.2024.00058

會議時間：2024.11.17-22

作者列表：

1） 盧玥辰 中國石油大學（北京）人工智能學院 先進科學與工程計算 博22

2） 曾禮杰 中國石油大學（北京）人工智能學院 計算機科學與技術 碩24

3） 王騰程 中國石油大學（北京）人工智能學院 計算機科學與技術 碩21

4） 付   旭 中國石油大學（北京）人工智能學院 計算機科學與技術 碩21

5） 李文瑄 中國石油大學（北京）人工智能學院 計算機科學與技術 碩24

6） 程賀琳 中國石油大學（北京）人工智能學院 計算機科學與技術 碩24

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

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

9） Marc Casas 巴塞羅那超級計算中心 首席研究員

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

摘要:

代數多重網格（AMG）方法因其靈活性和適應性，被廣泛用于求解稀疏線性系統。盡管現代GPU為AMG提供了大規模并行計算能力，但其最新硬件特性（即Tensor Core及低精度計算能力）尚未被很好地用于加速AMG。本文提出了一個新的AMG解法器AmgT，其能夠在AMG算法的多個階段中利用最新GPU的Tensor Core和混合精度能力。我們首先提出了一種統一稀疏存儲格式，充分利用Tensor Core及混合精度來提高AMG算法中頻繁調用的稀疏矩陣-矩陣乘法（SpGEMM）和稀疏矩陣-向量乘法（SpMV）的性能，并盡量減少AMG整個數據流中格式轉換的成本。同時，為更好地利用現有庫中的算法組件，AmgT的數據格式和計算kernel被集成到HYPRE庫中。實驗結果表明，在NVIDIA A100、H100 和AMD MI210 GPU上，AmgT比原GPU版本的HYPRE平均快1.46倍、1.32倍和2.24倍（最高達2.10倍、2.06倍和3.67倍）。

背景與動機:

代數多重網格（AMG）方法因其靈活性和適應性，被廣泛用于求解稀疏線性系統。盡管現代GPU為AMG提供了大規模并行計算能力，但其最新硬件特性（即Tensor Core及其低精度計算能力）尚未被很好的用于加速AMG。然而利用Tensor Core加速AMG面臨著三個挑戰：（1）存儲方面，如何避免為SpGEMM和SpMV生成不同格式的矩陣；（2）計算方面，如何將一般稀疏結構與Tensor Core的嚴格稠密GEMM模式相匹配；（3）精度方面，如何將可變精度的SpGEMM和SpMV集成到AMG的完整數據流中。

設計與實現:

我們首先設計了一種統一稀疏存儲格式mBSR（圖1），這是經典塊稀疏行（BSR）格式的一種變體。mBSR格式將稀疏矩陣存儲在一組大小為4×4的稠密塊中，并使用位圖存儲每個塊中的非零元位置。

 

圖1. mBSR格式 

在這種數據格式的基礎上，我們提出了一種新的AMG求解器AmgT，它在AMG算法的多個階段利用了最新GPU的Tensor Core和及其低精度能力。AmgT基于新的SpGEMM（圖2）和SpMV算法，能夠根據塊的稀疏程度同時使用Tensor Core和CUDA Core加速計算。AmgT中的SpGEMM算法會先分析矩陣數據，并將所有塊行歸入八個分區，然后執行兩步哈希操作進行符號計算，以獲取生成矩陣中塊的位置信息，最后同時使用Tensor Core和CUDA Core進行數值計算。AmgT中的SpMV算法采用自適應性選擇負載均衡和計算kernel策略，并同樣實現了Tensor Core和CUDA Core的混合使用，以提高整體性能。

 

圖2. mBSR格式的SpGEMM計算過程

最后，多種精度的新SpGEMM和SpMV算法分別在不同層的網格中調用，以有效利用Tensor Core的計算能力。

 

圖3. AmgT的完整數據流

實驗結果及分析:

為更好地利用現有庫中的算法組件，AmgT的數據格式和計算kernel被集成到HYPRE庫中。我們的測試平臺為兩款NVIDIA GPU A100（Ampere）和H100（Hopper）以及一款AMG GPU MI210（CDNA2），并測試了SuiteSparse矩陣集的16個代表性矩陣。

實驗結果表明，在NVIDIA A100、H100 和AMD MI210 GPU上，AmgT比原GPU版本的HYPRE平均快1.46倍、1.32倍和2.24倍（最高達2.10倍、2.06倍和3.67倍），如圖4所示。在A100和H100上，我們的混合精度AmgT比雙精度AmgT分別快1.03倍和1.04倍（最高1.08倍和1.14倍）。此外，獨立kernel測試表明，我們的SpGEMM比cuSPARSE和rocSPARSE SpGEMM分別快3.09倍、2.40倍和4.67倍（最高達7.61倍、6.11倍和5.96倍），SpMV在三個GPU上的性能分別是cuSPARSE和rocSPARSE SpMV的1.34倍、1.19倍和2.92倍（最高達2.21倍、2.09倍和6.70倍）。

圖4. 不同方法的AMG在三個GPU上的性能比較 

我們還比較了雙精度HYPRE、AmgT和混合精度AmgT在8個A100 GPU上的性能，如圖5所示。與雙精度HYPRE相比，我們的雙精度AmgT方法的速度快了1.35倍（最高1.84倍）。此外，與雙精度AmgT方法相比，我們的混合精度AmgT方法的速度快了1.06 倍（最高 1.11 倍）。雖然數據劃分會導致通信成本增加、每GPU的計算量減少，但我們的算法仍能保持相對于調用cuSPARSE 的HYPRE的穩定優勢。

 

圖5. HYPRE和AmgT在8個A100 GPU上的性能比較

作者簡介:

劉偉峰，教授，博士生導師，歐盟瑪麗居里學者。2002年和2006年于中國石油大學（北京）計算機系獲學士與碩士學位。2006年至2012年在中國石化石油勘探開發研究院歷任助理工程師、工程師和高級研究師，其間主要研究領域為石油地球物理勘探的高性能算法。2016年于丹麥哥本哈根大學獲計算科學博士學位，主要研究方向為數值線性代數和并行計算，其中尤其關注稀疏矩陣的數據結構、并行算法和軟件。研究工作發表于SC、PPoPP、DAC、ASPLOS、ICS、IPDPS、ICPP、TPDS、JPDC、FGCS和Parco等重要國際會議和期刊。擔任TPDS、SISC和TKDE等多個重要國際期刊審稿人，以及SC、ICS、IPDPS和ICPP等多個重要國際會議的程序委員會委員。他是IEEE高級會員、CCF高級會員、ACM和SIAM會員。

聯系方式：weifeng.liu@cup.edu.cn。